refactor file-stt restate pipeline

Author: yujonglee

apps/restate/src/audioPipeline.ts

@@ -0,0 +1,52 @@
+import { CallbackUrl, createClient } from "@deepgram/sdk";
+import { z } from "zod";
+
+export const DeepgramCallback = z.object({
+  results: z
+    .object({
+      channels: z
+        .array(
+          z.object({
+            alternatives: z
+              .array(z.object({ transcript: z.string() }))
+              .optional(),
+          }),
+        )
+        .optional(),
+    })
+    .optional(),
+  channel: z
+    .object({
+      alternatives: z.array(z.object({ transcript: z.string() })).optional(),
+    })
+    .optional(),
+});
+
+export type DeepgramCallbackType = z.infer<typeof DeepgramCallback>;
+
+export function extractTranscript(payload: DeepgramCallbackType): string {
+  return (
+    payload.results?.channels?.[0]?.alternatives?.[0]?.transcript ??
+    payload.channel?.alternatives?.[0]?.transcript ??
+    ""
+  );
+}
+
+export async function transcribeWithCallback(
+  audioUrl: string,
+  callbackUrl: string,
+  apiKey: string,
+): Promise<string> {
+  const client = createClient(apiKey);
+  const { result, error } =
+    await client.listen.prerecorded.transcribeUrlCallback(
+      { url: audioUrl },
+      new CallbackUrl(callbackUrl),
+      { model: "nova-3", smart_format: true },
+    );
+
+  if (error) throw new Error(`Deepgram: ${error.message}`);
+  if (!result?.request_id) throw new Error("Deepgram: missing request_id");
+
+  return result.request_id;
+}

apps/restate/src/deepgram.ts

@@ -1,6 +1,5 @@
 import * as restate from "@restatedev/restate-sdk-cloudflare-workers/fetch";
 
-import { audioPipeline } from "./audioPipeline";
 import { type Env, envSchema } from "./env";
 import { rateLimiter } from "./services/rate-limit";
 import { sttFile } from "./services/stt-file";
@@ -9,7 +8,7 @@ export default {
   fetch(request: Request, _env: Env, _ctx: ExecutionContext) {
     const env = envSchema.parse(_env);
     return restate.createEndpointHandler({
-      services: [audioPipeline, rateLimiter, sttFile],
+      services: [rateLimiter, sttFile],
       ...(env.RESTATE_IDENTITY_KEY
         ? { identityKeys: [env.RESTATE_IDENTITY_KEY] }
         : {}),

apps/restate/src/index.ts

@@ -1,15 +1,5 @@
 import * as restate from "@restatedev/restate-sdk-cloudflare-workers/fetch";
 
-interface RateLimitState {
-  windowStartMs: number;
-  count: number;
-}
-
-interface RateLimitConfig {
-  windowMs: number;
-  maxInWindow: number;
-}
-
 export const rateLimiter = restate.object({
   name: "RateLimiter",
   handlers: {
@@ -28,7 +18,7 @@ export const rateLimiter = restate.object({
 
       if (current.count >= config.maxInWindow) {
         ctx.set("state", current);
-        throw new restate.TerminalError("Rate limit exceeded", {
+        throw new restate.TerminalError("rate_limit_exceeded", {
           errorCode: 429,
         });
       }
@@ -47,3 +37,13 @@ export type RateLimiter = typeof rateLimiter;
 export function limiter(ctx: restate.Context, key: string) {
   return ctx.objectClient<RateLimiter>(rateLimiter, key);
 }
+
+interface RateLimitState {
+  windowStartMs: number;
+  count: number;
+}
+
+interface RateLimitConfig {
+  windowMs: number;
+  maxInWindow: number;
+}

apps/restate/src/services/rate-limit.ts

@@ -1,8 +1,13 @@
-import { CallbackUrl, createClient } from "@deepgram/sdk";
 import * as restate from "@restatedev/restate-sdk-cloudflare-workers/fetch";
 import { serde } from "@restatedev/restate-sdk-zod";
 import { z } from "zod";
 
+import {
+  DeepgramCallback,
+  DeepgramCallbackType,
+  extractTranscript,
+  transcribeWithCallback,
+} from "../deepgram";
 import { type Env } from "../env";
 import { createSignedUrl, deleteFile } from "../supabase";
 import { limiter } from "./rate-limit";
@@ -18,63 +23,12 @@ const SttStatus = z.enum(["QUEUED", "TRANSCRIBING", "DONE", "ERROR"]);
 
 export type SttStatusType = z.infer<typeof SttStatus>;
 
-const DeepgramCallback = z.object({
-  results: z
-    .object({
-      channels: z
-        .array(
-          z.object({
-            alternatives: z
-              .array(z.object({ transcript: z.string() }))
-              .optional(),
-          }),
-        )
-        .optional(),
-    })
-    .optional(),
-  channel: z
-    .object({
-      alternatives: z.array(z.object({ transcript: z.string() })).optional(),
-    })
-    .optional(),
-});
-
-export type DeepgramCallbackType = z.infer<typeof DeepgramCallback>;
-
-interface SttResult {
-  status: SttStatusType;
-  transcript?: string;
-  error?: string;
-}
-
-async function transcribe(
-  audioUrl: string,
-  callbackUrl: string,
-  apiKey: string,
-): Promise<string> {
-  const client = createClient(apiKey);
-  const { result, error } =
-    await client.listen.prerecorded.transcribeUrlCallback(
-      { url: audioUrl },
-      new CallbackUrl(callbackUrl),
-      { model: "nova-3", smart_format: true },
-    );
-
-  if (error) throw new Error(`Deepgram: ${error.message}`);
-  if (!result?.request_id) throw new Error("Deepgram: missing request_id");
-
-  return result.request_id;
-}
-
 export const sttFile = restate.workflow({
   name: "SttFile",
   handlers: {
     run: restate.handlers.workflow.workflow(
       { input: serde.zod(SttFileInput) },
-      async (
-        ctx: restate.WorkflowContext,
-        input: SttFileInputType,
-      ): Promise<SttResult> => {
+      async (ctx: restate.WorkflowContext, input: SttFileInputType) => {
         ctx.set("status", "QUEUED" as SttStatusType);
         ctx.set("fileId", input.fileId);
 
@@ -95,15 +49,15 @@ export const sttFile = restate.workflow({
           const callbackUrl = `${env.RESTATE_INGRESS_URL.replace(/\/+$/, "")}/SttFile/${encodeURIComponent(ctx.key)}/onTranscript`;
 
           const requestId = await ctx.run("transcribe", () =>
-            transcribe(audioUrl, callbackUrl, env.DEEPGRAM_API_KEY),
+            transcribeWithCallback(audioUrl, callbackUrl, env.DEEPGRAM_API_KEY),
           );
           ctx.set("deepgramRequestId", requestId);
 
           const transcript = await ctx.promise<string>("transcript");
           ctx.set("transcript", transcript);
           ctx.set("status", "DONE" as SttStatusType);
 
-          return { status: "DONE", transcript };
+          return { status: "DONE" as const, transcript };
         } catch (err) {
           const error = err instanceof Error ? err.message : "Unknown error";
           ctx.set("status", "ERROR" as SttStatusType);
@@ -126,18 +80,13 @@ export const sttFile = restate.workflow({
         const existing = await ctx.get<string>("transcript");
         if (existing !== undefined) return;
 
-        const transcript =
-          payload.results?.channels?.[0]?.alternatives?.[0]?.transcript ??
-          payload.channel?.alternatives?.[0]?.transcript ??
-          "";
-
-        ctx.promise<string>("transcript").resolve(transcript);
+        ctx.promise<string>("transcript").resolve(extractTranscript(payload));
       },
     ),
 
     getStatus: restate.handlers.workflow.shared(
       {},
-      async (ctx: restate.WorkflowSharedContext): Promise<SttResult> => {
+      async (ctx: restate.WorkflowSharedContext) => {
         const status = (await ctx.get<SttStatusType>("status")) ?? "QUEUED";
         const transcript = await ctx.get<string>("transcript");
         const error = await ctx.get<string>("error");