diff --git a/packages/media-understanding-common/src/format.ts b/packages/media-understanding-common/src/format.ts index 2fd0249dd136..b931a71cd663 100644 --- a/packages/media-understanding-common/src/format.ts +++ b/packages/media-understanding-common/src/format.ts @@ -1,9 +1,18 @@ // Media Understanding Common helper module supports format behavior. import type { MediaUnderstandingOutput } from "./types.js"; +const sectionByKind = { + "audio.transcription": { title: "Audio", label: "Transcript" }, + "image.description": { title: "Image", label: "Description" }, + "video.description": { title: "Video", label: "Description" }, +} satisfies Record< + MediaUnderstandingOutput["kind"], + { title: string; label: "Transcript" | "Description" } +>; + function formatSection( title: string, - kind: "Transcript" | "Description", + label: "Transcript" | "Description", text: string, userText?: string, ): string { @@ -11,7 +20,7 @@ function formatSection( if (userText) { lines.push(`User text:\n${userText}`); } - lines.push(`${kind}:\n${text}`); + lines.push(`${label}:\n${text}`); return lines.join("\n"); } @@ -42,32 +51,11 @@ export function formatMediaUnderstandingBody(params: { const next = (seen.get(output.kind) ?? 0) + 1; seen.set(output.kind, next); const suffix = count > 1 ? ` ${next}/${count}` : ""; - if (output.kind === "audio.transcription") { - sections.push( - formatSection( - `Audio${suffix}`, - "Transcript", - output.text, - outputs.length === 1 ? userText : undefined, - ), - ); - continue; - } - if (output.kind === "image.description") { - sections.push( - formatSection( - `Image${suffix}`, - "Description", - output.text, - outputs.length === 1 ? userText : undefined, - ), - ); - continue; - } + const section = sectionByKind[output.kind]; sections.push( formatSection( - `Video${suffix}`, - "Description", + `${section.title}${suffix}`, + section.label, output.text, outputs.length === 1 ? userText : undefined, ),