@markusylisiurunen/tau 0.3.64 → 0.3.65

package/dist/core/utils/gemini_transcription.js

@@ -1,8 +1,20 @@
+import WebSocket from "ws";
 import { z } from "zod";
 import { formatSpeechToTextContext } from "./speech_to_text_context.js";
+import { normalizeSpeechToTextKeywords, SPEECH_TO_TEXT_KEYWORD_INSTRUCTIONS, } from "./speech_to_text_keywords.js";
 const GEMINI_GENERATE_CONTENT_BASE_URL = "https://generativelanguage.googleapis.com/v1beta/models";
-const GEMINI_TRANSCRIPTION_MODEL = "gemini-3.7-flash";
-const GEMINI_TRANSCRIPTION_THINKING_LEVEL = "low";
+const GEMINI_FILE_UPLOAD_URL = "https://generativelanguage.googleapis.com/upload/v1beta/files";
+const GEMINI_INTERACTIONS_URL = "https://generativelanguage.googleapis.com/v1beta/interactions";
+const GEMINI_FILES_BASE_URL = "https://generativelanguage.googleapis.com/v1beta";
+const GEMINI_LIVE_TRANSCRIPTION_URL = "wss://generativelanguage.googleapis.com/ws/google.ai.generativelanguage.v1beta.GenerativeService.BidiGenerateContent";
+const GEMINI_TRANSCRIPTION_MODEL = "gemini-3.5-transcribe";
+const GEMINI_LIVE_TRANSCRIPTION_MODEL = "gemini-3.5-transcribe-live";
+const GEMINI_TRANSCRIPTION_KEYWORD_MODEL = "gemini-3.7-flash";
+const GEMINI_TRANSCRIPTION_MAX_KEYWORD_CHARACTERS_TOTAL = 10_000;
+const GEMINI_TRANSCRIPTION_KEYWORD_TIMEOUT_MS = 15_000;
+const GEMINI_TRANSCRIPTION_CONNECT_TIMEOUT_MS = 15_000;
+const GEMINI_TRANSCRIPTION_COMPLETION_TIMEOUT_MS = 30_000;
+const GEMINI_FILE_DELETE_TIMEOUT_MS = 5_000;
 const DEFAULT_GEMINI_AUDIO_MIME_TYPE = "audio/wav";
 const errorPayloadSchema = z.object({
     error: z
@@ -11,15 +23,7 @@ const errorPayloadSchema = z.object({
     })
         .optional(),
 });
-const GEMINI_TRANSCRIPTION_RESPONSE_SCHEMA = {
-    type: "OBJECT",
-    properties: {
-        transcription: { type: "STRING" },
-    },
-    required: ["transcription"],
-};
-const textPartSchema = z.object({ text: z.string() });
-const apiResponseSchema = z.object({
+const generateContentResponseSchema = z.object({
     candidates: z
         .array(z.object({
         content: z.object({
@@ -28,122 +32,531 @@ const apiResponseSchema = z.object({
     }))
         .optional(),
 });
-const transcriptionResultSchema = z.object({
-    transcription: z.string().trim().min(1),
+const textPartSchema = z.object({ text: z.string() });
+const transcriptionKeywordsSchema = z
+    .object({
+    keywords: z.array(z.string()),
+})
+    .strict();
+const uploadedFileSchema = z.object({
+    file: z.object({
+        name: z.string().trim().min(1),
+        uri: z.string().trim().min(1),
+    }),
+});
+const interactionSchema = z.object({
+    steps: z.array(z.unknown()).optional(),
 });
+const interactionModelOutputSchema = z.object({
+    type: z.literal("model_output"),
+    content: z.array(z.unknown()),
+});
+const interactionTextSchema = z.object({
+    type: z.literal("text"),
+    text: z.string(),
+});
+const liveMessageSchema = z
+    .object({
+    setupComplete: z.object({}).optional(),
+    serverContent: z
+        .object({
+        inputTranscription: z
+            .object({
+            text: z.string().trim().min(1),
+        })
+            .optional(),
+    })
+        .passthrough()
+        .optional(),
+    error: z
+        .object({
+        message: z.string().trim().min(1),
+    })
+        .optional(),
+})
+    .passthrough();
+class GeminiStreamingTranscriptionImpl {
+    socket;
+    keywordAbortController = new AbortController();
+    keywordsPromise;
+    ready = false;
+    aborted = false;
+    failure;
+    completedTranscript;
+    hasAudio = false;
+    pendingAudio = [];
+    readyTimeout;
+    completionTimeout;
+    resolveReady;
+    rejectReady;
+    resolveCompletion;
+    rejectCompletion;
+    readyPromise;
+    constructor(options) {
+        const apiKey = options.apiKey.trim();
+        if (!apiKey) {
+            throw new Error("missing Gemini API key");
+        }
+        this.keywordsPromise = prepareGeminiTranscriptionKeywords({
+            apiKey,
+            context: options.context,
+            signal: this.keywordAbortController.signal,
+            fetchImpl: options.fetchImpl,
+        });
+        const webSocketFactory = options.webSocketFactory ?? defaultWebSocketFactory;
+        this.socket = webSocketFactory(`${GEMINI_LIVE_TRANSCRIPTION_URL}?key=${encodeURIComponent(apiKey)}`);
+        this.readyPromise = new Promise((resolve, reject) => {
+            this.resolveReady = resolve;
+            this.rejectReady = reject;
+        });
+        void this.readyPromise.catch(() => { });
+        this.readyTimeout = setTimeout(() => {
+            this.fail(new Error("timed out opening Gemini transcription session"));
+            this.socket.terminate();
+        }, GEMINI_TRANSCRIPTION_CONNECT_TIMEOUT_MS);
+        this.readyTimeout.unref?.();
+        this.socket.on("open", () => {
+            if (this.readyTimeout)
+                clearTimeout(this.readyTimeout);
+            this.readyTimeout = undefined;
+            void this.configureSession();
+        });
+        this.socket.on("message", (data) => this.handleMessage(data));
+        this.socket.on("error", (error) => {
+            this.fail(new Error(`Gemini transcription connection failed: ${error.message}`));
+        });
+        this.socket.on("close", (code, reason) => {
+            if (this.aborted || this.completedTranscript)
+                return;
+            const detail = reason.toString("utf8").trim();
+            this.fail(new Error(detail
+                ? `Gemini transcription connection closed (${code}): ${detail}`
+                : `Gemini transcription connection closed (${code})`));
+        });
+    }
+    appendAudio(audio) {
+        if (audio.length === 0 || this.aborted || this.failure || this.completedTranscript)
+            return;
+        this.hasAudio = true;
+        if (!this.ready) {
+            this.pendingAudio.push(audio);
+            return;
+        }
+        this.sendAudio(audio);
+    }
+    async finish(options = {}) {
+        const abortListener = () => this.abort();
+        if (options.signal?.aborted) {
+            abortListener();
+        }
+        else {
+            options.signal?.addEventListener("abort", abortListener, { once: true });
+        }
+        try {
+            await this.readyPromise;
+            if (this.failure)
+                throw this.failure;
+            if (this.aborted)
+                throw new Error("Gemini transcription was aborted");
+            if (this.completedTranscript)
+                return this.completedTranscript;
+            if (!this.hasAudio)
+                throw new Error("Gemini transcription received no audio");
+            const completion = new Promise((resolve, reject) => {
+                this.resolveCompletion = resolve;
+                this.rejectCompletion = reject;
+            });
+            this.completionTimeout = setTimeout(() => {
+                this.fail(new Error("timed out waiting for Gemini transcription"));
+                this.socket.terminate();
+            }, GEMINI_TRANSCRIPTION_COMPLETION_TIMEOUT_MS);
+            this.completionTimeout.unref?.();
+            this.send({ realtimeInput: { activityEnd: {} } });
+            return await completion;
+        }
+        finally {
+            options.signal?.removeEventListener("abort", abortListener);
+        }
+    }
+    abort() {
+        if (this.aborted || this.completedTranscript)
+            return;
+        const error = new Error("Gemini transcription was aborted");
+        this.aborted = true;
+        this.pendingAudio = [];
+        this.keywordAbortController.abort(error);
+        this.clearTimers();
+        this.rejectReady?.(error);
+        this.rejectCompletion?.(error);
+        this.clearWaiters();
+        this.socket.close();
+    }
+    async configureSession() {
+        const keywords = await this.keywordsPromise;
+        if (this.aborted || this.failure)
+            return;
+        this.readyTimeout = setTimeout(() => {
+            this.fail(new Error("timed out opening Gemini transcription session"));
+            this.socket.terminate();
+        }, GEMINI_TRANSCRIPTION_CONNECT_TIMEOUT_MS);
+        this.readyTimeout.unref?.();
+        this.send({
+            setup: {
+                model: `models/${GEMINI_LIVE_TRANSCRIPTION_MODEL}`,
+                generationConfig: {
+                    responseModalities: ["TEXT"],
+                },
+                inputAudioTranscription: {
+                    languageCodes: [],
+                    ...(keywords.length > 0 ? { customVocabulary: keywords } : {}),
+                    mode: "SMART",
+                },
+                realtimeInputConfig: {
+                    automaticActivityDetection: {
+                        disabled: true,
+                    },
+                },
+            },
+        });
+    }
+    handleMessage(data) {
+        let payload;
+        try {
+            payload = JSON.parse(formatWebSocketMessage(data));
+        }
+        catch {
+            this.fail(new Error("Gemini transcription returned malformed JSON"));
+            return;
+        }
+        const event = liveMessageSchema.safeParse(payload);
+        if (!event.success) {
+            this.fail(new Error("Gemini transcription returned a malformed event"));
+            return;
+        }
+        if (event.data.error) {
+            this.fail(new Error(event.data.error.message));
+            return;
+        }
+        if (event.data.setupComplete && !this.ready) {
+            this.ready = true;
+            if (this.readyTimeout)
+                clearTimeout(this.readyTimeout);
+            this.readyTimeout = undefined;
+            this.send({ realtimeInput: { activityStart: {} } });
+            const pendingAudio = this.pendingAudio;
+            this.pendingAudio = [];
+            for (const audio of pendingAudio) {
+                this.sendAudio(audio);
+            }
+            if (this.failure)
+                return;
+            this.resolveReady?.();
+            this.resolveReady = undefined;
+            this.rejectReady = undefined;
+            return;
+        }
+        const transcript = event.data.serverContent?.inputTranscription?.text;
+        if (transcript) {
+            this.completedTranscript = transcript;
+            this.clearTimers();
+            this.resolveCompletion?.(transcript);
+            this.clearWaiters();
+            this.socket.close();
+        }
+    }
+    sendAudio(audio) {
+        this.send({
+            realtimeInput: {
+                audio: {
+                    data: audio.toString("base64"),
+                    mimeType: "audio/pcm;rate=16000",
+                },
+            },
+        });
+    }
+    send(event) {
+        if (this.failure || this.aborted)
+            return;
+        try {
+            this.socket.send(JSON.stringify(event), (error) => {
+                if (error) {
+                    this.fail(new Error(`failed to send Gemini transcription audio: ${error.message}`));
+                }
+            });
+        }
+        catch (error) {
+            this.fail(new Error(`failed to send Gemini transcription audio: ${error.message}`));
+        }
+    }
+    fail(error) {
+        if (this.failure || this.aborted || this.completedTranscript)
+            return;
+        this.failure = error;
+        this.pendingAudio = [];
+        this.keywordAbortController.abort(error);
+        this.clearTimers();
+        this.rejectReady?.(error);
+        this.rejectCompletion?.(error);
+        this.clearWaiters();
+        this.socket.terminate();
+    }
+    clearTimers() {
+        if (this.readyTimeout)
+            clearTimeout(this.readyTimeout);
+        if (this.completionTimeout)
+            clearTimeout(this.completionTimeout);
+        this.readyTimeout = undefined;
+        this.completionTimeout = undefined;
+    }
+    clearWaiters() {
+        this.resolveReady = undefined;
+        this.rejectReady = undefined;
+        this.resolveCompletion = undefined;
+        this.rejectCompletion = undefined;
+    }
+}
+export function startGeminiTranscription(options) {
+    return new GeminiStreamingTranscriptionImpl(options);
+}
 export async function transcribeGeminiAudio(options) {
     const apiKey = options.apiKey.trim();
     if (!apiKey) {
         throw new Error("missing Gemini API key");
     }
     const fetchFn = options.fetchImpl ?? fetch;
-    const response = await fetchFn(`${GEMINI_GENERATE_CONTENT_BASE_URL}/${GEMINI_TRANSCRIPTION_MODEL}:generateContent`, {
-        method: "POST",
-        headers: {
-            "Content-Type": "application/json",
-            "x-goog-api-key": apiKey,
-        },
-        signal: options.signal,
-        body: JSON.stringify({
-            systemInstruction: {
-                parts: [
+    const [keywords, uploadedFile] = await Promise.all([
+        prepareGeminiTranscriptionKeywords({
+            apiKey,
+            context: options.context,
+            signal: options.signal,
+            fetchImpl: fetchFn,
+        }),
+        uploadGeminiAudio({
+            apiKey,
+            audio: options.audio,
+            mimeType: options.mimeType ?? DEFAULT_GEMINI_AUDIO_MIME_TYPE,
+            signal: options.signal,
+            fetchImpl: fetchFn,
+        }),
+    ]);
+    try {
+        const response = await fetchFn(GEMINI_INTERACTIONS_URL, {
+            method: "POST",
+            headers: {
+                "Content-Type": "application/json",
+                "x-goog-api-key": apiKey,
+            },
+            signal: options.signal,
+            body: JSON.stringify({
+                model: GEMINI_TRANSCRIPTION_MODEL,
+                input: [
                     {
-                        text: buildTranscriptionSystemInstruction(),
+                        type: "audio",
+                        uri: uploadedFile.uri,
+                        mime_type: options.mimeType ?? DEFAULT_GEMINI_AUDIO_MIME_TYPE,
                     },
                 ],
+                generation_config: {
+                    transcription_config: {
+                        language_codes: [],
+                        ...(keywords.length > 0 ? { custom_vocabulary: keywords } : {}),
+                        mode: "smart",
+                    },
+                },
+                store: false,
+            }),
+        });
+        const { payload, responseText } = await readResponsePayload(response);
+        if (!response.ok) {
+            throw createResponseError(payload, responseText, response.status);
+        }
+        const transcript = extractInteractionText(payload)?.trim();
+        if (!transcript) {
+            throw new Error("transcription result was empty or malformed");
+        }
+        return transcript;
+    }
+    finally {
+        await deleteGeminiFile({
+            apiKey,
+            fileName: uploadedFile.name,
+            fetchImpl: fetchFn,
+        });
+    }
+}
+async function prepareGeminiTranscriptionKeywords(args) {
+    const formattedContext = formatSpeechToTextContext(args.context);
+    if (!formattedContext)
+        return [];
+    try {
+        const signals = [AbortSignal.timeout(GEMINI_TRANSCRIPTION_KEYWORD_TIMEOUT_MS)];
+        if (args.signal)
+            signals.push(args.signal);
+        const response = await (args.fetchImpl ?? fetch)(`${GEMINI_GENERATE_CONTENT_BASE_URL}/${GEMINI_TRANSCRIPTION_KEYWORD_MODEL}:generateContent`, {
+            method: "POST",
+            headers: {
+                "Content-Type": "application/json",
+                "x-goog-api-key": args.apiKey,
             },
-            contents: [
-                {
-                    parts: [
-                        {
-                            text: buildTranscriptionPrompt(options.context),
-                        },
-                        {
-                            inlineData: {
-                                mimeType: options.mimeType ?? DEFAULT_GEMINI_AUDIO_MIME_TYPE,
-                                data: options.audio.toString("base64"),
+            signal: AbortSignal.any(signals),
+            body: JSON.stringify({
+                systemInstruction: {
+                    parts: [{ text: SPEECH_TO_TEXT_KEYWORD_INSTRUCTIONS }],
+                },
+                contents: [
+                    {
+                        parts: [{ text: formattedContext }],
+                    },
+                ],
+                generationConfig: {
+                    responseMimeType: "application/json",
+                    responseSchema: {
+                        type: "OBJECT",
+                        properties: {
+                            keywords: {
+                                type: "ARRAY",
+                                items: { type: "STRING" },
                             },
                         },
-                    ],
-                },
-            ],
-            generationConfig: {
-                responseMimeType: "application/json",
-                responseSchema: GEMINI_TRANSCRIPTION_RESPONSE_SCHEMA,
-                thinkingConfig: {
-                    thinkingLevel: GEMINI_TRANSCRIPTION_THINKING_LEVEL,
+                        required: ["keywords"],
+                    },
+                    thinkingConfig: {
+                        thinkingLevel: "low",
+                    },
                 },
-            },
-        }),
-    });
-    const responseText = await response.text();
-    let payload;
-    try {
-        payload = responseText ? JSON.parse(responseText) : undefined;
+            }),
+        });
+        if (!response.ok) {
+            await response.body?.cancel().catch(() => { });
+            return [];
+        }
+        const payload = (await response.json());
+        const outputText = extractGenerateContentText(payload);
+        const parsedKeywords = transcriptionKeywordsSchema.safeParse(outputText ? JSON.parse(outputText) : undefined);
+        return parsedKeywords.success
+            ? normalizeSpeechToTextKeywords(parsedKeywords.data.keywords, {
+                maxTotalCharacters: GEMINI_TRANSCRIPTION_MAX_KEYWORD_CHARACTERS_TOTAL,
+            })
+            : [];
     }
     catch {
-        payload = undefined;
+        return [];
     }
-    if (!response.ok) {
-        const parsed = errorPayloadSchema.safeParse(payload);
-        const fallbackMessage = responseText.trim() || `HTTP ${response.status}`;
-        throw new Error(parsed.success ? (parsed.data.error?.message ?? fallbackMessage) : fallbackMessage);
+}
+async function uploadGeminiAudio(args) {
+    const startResponse = await args.fetchImpl(GEMINI_FILE_UPLOAD_URL, {
+        method: "POST",
+        headers: {
+            "Content-Type": "application/json",
+            "x-goog-api-key": args.apiKey,
+            "X-Goog-Upload-Protocol": "resumable",
+            "X-Goog-Upload-Command": "start",
+            "X-Goog-Upload-Header-Content-Length": String(args.audio.byteLength),
+            "X-Goog-Upload-Header-Content-Type": args.mimeType,
+        },
+        signal: args.signal,
+        body: JSON.stringify({ file: { display_name: "tau-speech" } }),
+    });
+    if (!startResponse.ok) {
+        const { payload, responseText } = await readResponsePayload(startResponse);
+        throw createResponseError(payload, responseText, startResponse.status);
     }
-    const text = extractGeminiText(payload);
-    if (!text) {
-        throw new Error("transcription result was empty or malformed");
+    const uploadUrl = startResponse.headers.get("x-goog-upload-url")?.trim();
+    await startResponse.body?.cancel().catch(() => { });
+    if (!uploadUrl) {
+        throw new Error("Gemini file upload did not return an upload URL");
     }
-    return text;
-}
-function buildTranscriptionSystemInstruction() {
-    return [
-        "You are a speech-to-text engine.",
-        "Transcribe the speaker's intended message for insertion into a chat input.",
-        "Detect the speaker's language and transcribe in that same language; never translate unless the speaker explicitly asks for translation.",
-        "Preserve the speaker's wording and register as spoken, including colloquial forms, dialect, and informal language; do not normalize informal speech into formal standard language.",
-        "Use natural punctuation and capitalization where helpful, without changing the speaker's wording or register.",
-        "Lightly clean only speech artifacts that do not affect meaning, such as filler words, repeated stutters, obvious false starts, and unintelligible mumbling.",
-        "Do not rewrite, paraphrase, summarize, answer the speaker, add labels, add timestamps, or describe background sounds.",
-    ].join("\n");
+    const uploadResponse = await args.fetchImpl(uploadUrl, {
+        method: "POST",
+        headers: {
+            "Content-Length": String(args.audio.byteLength),
+            "Content-Type": args.mimeType,
+            "X-Goog-Upload-Offset": "0",
+            "X-Goog-Upload-Command": "upload, finalize",
+        },
+        signal: args.signal,
+        body: Uint8Array.from(args.audio),
+    });
+    const { payload, responseText } = await readResponsePayload(uploadResponse);
+    if (!uploadResponse.ok) {
+        throw createResponseError(payload, responseText, uploadResponse.status);
+    }
+    const uploadedFile = uploadedFileSchema.safeParse(payload);
+    if (!uploadedFile.success) {
+        throw new Error("Gemini file upload returned a malformed response");
+    }
+    return uploadedFile.data.file;
 }
-function buildTranscriptionPrompt(context) {
-    const formattedContext = formatSpeechToTextContext(context);
-    return [
-        "Transcribe the attached audio into the transcription field.",
-        "Return only the lightly cleaned transcript text, with no timestamps or commentary.",
-        formattedContext
-            ? [
-                "Use the recent conversation context below only to resolve likely words, names, acronyms, terminology, and references in the audio.",
-                "Do not transcribe the context itself, and do not add words from the context that were not spoken.",
-                "If the audio is ambiguous, prefer the transcription that best fits this context.",
-                "",
-                "Recent conversation context:",
-                formattedContext,
-            ].join("\n")
-            : undefined,
-    ]
-        .filter((line) => line !== undefined)
-        .join("\n");
+async function deleteGeminiFile(args) {
+    try {
+        const response = await args.fetchImpl(`${GEMINI_FILES_BASE_URL}/${args.fileName}`, {
+            method: "DELETE",
+            headers: { "x-goog-api-key": args.apiKey },
+            signal: AbortSignal.timeout(GEMINI_FILE_DELETE_TIMEOUT_MS),
+        });
+        await response.body?.cancel().catch(() => { });
+    }
+    catch {
+        // best-effort remote cleanup
+    }
 }
-function extractGeminiText(payload) {
-    const parsed = apiResponseSchema.safeParse(payload);
-    if (!parsed.success) {
+function extractGenerateContentText(payload) {
+    const parsed = generateContentResponseSchema.safeParse(payload);
+    if (!parsed.success)
         return undefined;
-    }
-    const responseText = (parsed.data.candidates?.[0]?.content.parts ?? [])
+    return (parsed.data.candidates?.[0]?.content.parts ?? [])
         .map((part) => {
         const parsedPart = textPartSchema.safeParse(part);
         return parsedPart.success ? parsedPart.data.text : "";
     })
         .join("");
-    let transcriptionPayload;
+}
+function extractInteractionText(payload) {
+    const parsed = interactionSchema.safeParse(payload);
+    if (!parsed.success)
+        return undefined;
+    return (parsed.data.steps ?? [])
+        .flatMap((step) => {
+        const modelOutput = interactionModelOutputSchema.safeParse(step);
+        return modelOutput.success ? modelOutput.data.content : [];
+    })
+        .flatMap((content) => {
+        const text = interactionTextSchema.safeParse(content);
+        return text.success ? [text.data.text] : [];
+    })
+        .join("");
+}
+async function readResponsePayload(response) {
+    const responseText = await response.text();
     try {
-        transcriptionPayload = responseText ? JSON.parse(responseText) : undefined;
+        return {
+            payload: responseText ? JSON.parse(responseText) : undefined,
+            responseText,
+        };
     }
     catch {
-        return undefined;
+        return { payload: undefined, responseText };
     }
-    const transcription = transcriptionResultSchema.safeParse(transcriptionPayload);
-    return transcription.success ? transcription.data.transcription : undefined;
+}
+function createResponseError(payload, responseText, status) {
+    const parsed = errorPayloadSchema.safeParse(payload);
+    return new Error(parsed.success
+        ? (parsed.data.error?.message ?? (responseText.trim() || `HTTP ${status}`))
+        : responseText.trim() || `HTTP ${status}`);
+}
+function formatWebSocketMessage(data) {
+    if (typeof data === "string")
+        return data;
+    if (Buffer.isBuffer(data))
+        return data.toString("utf8");
+    if (data instanceof ArrayBuffer)
+        return Buffer.from(data).toString("utf8");
+    if (ArrayBuffer.isView(data)) {
+        return Buffer.from(data.buffer, data.byteOffset, data.byteLength).toString("utf8");
+    }
+    throw new Error("unsupported WebSocket message");
+}
+function defaultWebSocketFactory(url) {
+    return new WebSocket(url);
 }
 //# sourceMappingURL=gemini_transcription.js.map