Refactor translator service to be less jank and more modular

src/pages/developer/TestView.tsx

@@ -1,12 +1,83 @@
-import { useState } from "react";
+import { useCallback, useState } from "react";
 
 import { Button } from "@/components/buttons/Button";
+import { usePlayer } from "@/components/player/hooks/usePlayer";
+import { PlaybackErrorPart } from "@/pages/parts/player/PlaybackErrorPart";
+import { PlayerPart } from "@/pages/parts/player/PlayerPart";
+import {
+  CaptionListItem,
+  PlayerMeta,
+  playerStatus,
+} from "@/stores/player/slices/source";
+import { SourceSliceSource } from "@/stores/player/utils/qualities";
+
+const subtitlesTestMeta: PlayerMeta = {
+  type: "movie",
+  title: "Subtitles Test",
+  releaseYear: 2024,
+  tmdbId: "0",
+};
+
+const subtitlesTestSource: SourceSliceSource = {
+  type: "hls",
+  url: "http://localhost:8000/media/master.m3u8",
+};
+
+const subtitlesTestSubs: CaptionListItem[] = [
+  {
+    id: "http://localhost:8000/subs/en.srt",
+    display: "English",
+    language: "en",
+    url: "http://localhost:8000/subs/en.srt",
+    needsProxy: false,
+  },
+  {
+    id: "http://localhost:8000/subs/en-small.srt",
+    display: "English Small",
+    language: "en",
+    url: "http://localhost:8000/subs/en-small.srt",
+    needsProxy: false,
+  },
+  {
+    id: "http://localhost:8000/subs/ro.srt",
+    display: "Romanian",
+    language: "ro",
+    url: "http://localhost:8000/subs/ro.srt",
+    needsProxy: false,
+  },
+];
 
 // mostly empty view, add whatever you need
 export default function TestView() {
+  const player = usePlayer();
+  const [showPlayer, setShowPlayer] = useState(false);
   const [shouldCrash, setShouldCrash] = useState(false);
+
   if (shouldCrash) {
     throw new Error("I crashed");
   }
-  return <Button onClick={() => setShouldCrash(true)}>Crash me!</Button>;
+
+  const subtitlesTest = useCallback(async () => {
+    setShowPlayer(true);
+    player.reset();
+    await new Promise((r) => {
+      setTimeout(r, 100);
+    });
+    player.setShouldStartFromBeginning(true);
+    player.setMeta(subtitlesTestMeta);
+    player.playMedia(subtitlesTestSource, subtitlesTestSubs, null);
+  }, [player]);
+
+  return showPlayer ? (
+    <PlayerPart backUrl="/dev/">
+      {player && (player as any).status === playerStatus.PLAYBACK_ERROR ? (
+        <PlaybackErrorPart />
+      ) : null}
+    </PlayerPart>
+  ) : (
+    <>
+      <Button onClick={() => setShouldCrash(true)}>Crash me!</Button>
+      <Button onClick={() => subtitlesTest()}>Subtitles test</Button>
+    </>
+  );
 }

src/stores/player/slices/source.ts

@@ -9,10 +9,10 @@ import {
   selectQuality,
 } from "@/stores/player/utils/qualities";
 import { useQualityStore } from "@/stores/quality";
+import googletranslate from "@/utils/translation/googletranslate";
+import { translate } from "@/utils/translation/index";
 import { ValuesOf } from "@/utils/typeguard";
 
-import { translateSubtitle } from "../utils/captionstranslation";
-
 export const playerStatus = {
   IDLE: "idle",
   RESUME: "resume",
@@ -515,10 +515,10 @@ export const createSourceSlice: MakeSlice<SourceSlice> = (set, get) => ({
     }
 
     try {
-      const result = await translateSubtitle(
+      const result = await translate(
-        targetCaption.id,
+        store.caption.translateTask!.fetchedTargetCaption!,
-        store.caption.translateTask!.fetchedTargetCaption!.srtData,
         targetLanguage,
+        googletranslate,
       );
       if (cancelled) {
         return;

src/stores/player/utils/captionstranslation.ts

@@ -1,217 +0,0 @@
-import subsrt from "subsrt-ts";
-import { Caption, ContentCaption } from "subsrt-ts/dist/types/handler";
-
-const API_URL =
-  "https://translate.googleapis.com/translate_a/single?client=gtx&dt=t&dj=1&ie=UTF-8&oe=UTF-8&sl=auto";
-const RETRY_COUNT = 3;
-const FETCH_RATE = 100;
-const SUBTITLES_CACHE: Map<string, ArrayBuffer> = new Map<
-  string,
-  ArrayBuffer
->();
-
-async function compressStr(string: string): Promise<ArrayBuffer> {
-  const byteArray = new TextEncoder().encode(string);
-  const cs = new CompressionStream("deflate");
-  const writer = cs.writable.getWriter();
-  writer.write(byteArray);
-  writer.close();
-  return new Response(cs.readable).arrayBuffer();
-}
-
-async function decompressStr(byteArray: ArrayBuffer): Promise<string> {
-  const cs = new DecompressionStream("deflate");
-  const writer = cs.writable.getWriter();
-  writer.write(byteArray);
-  writer.close();
-  return new Response(cs.readable).arrayBuffer().then((arrayBuffer) => {
-    return new TextDecoder().decode(arrayBuffer);
-  });
-}
-
-function tryUseCachedCaption(
-  caption: ContentCaption,
-  cache: Map<string, string>,
-): boolean {
-  const text: string | undefined = cache.get(caption.text);
-  if (text) {
-    caption.text = text;
-    return true;
-  }
-  return false;
-}
-
-async function translateText(
-  text: string,
-  targetLang: string,
-): Promise<string | undefined> {
-  if (!text) {
-    return "";
-  }
-
-  const response = await (
-    await fetch(`${API_URL}&tl=${targetLang}&q=${encodeURIComponent(text)}`, {
-      method: "GET",
-      headers: {
-        Accept: "application/json",
-      },
-    })
-  ).json();
-
-  if (!response) {
-    throw new Error("Empty response");
-  }
-
-  return (response.sentences as any[])
-    .map((s: any) => s.trans as string)
-    .join("");
-}
-
-async function translateCaption(
-  caption: ContentCaption,
-  targetLang: string,
-): Promise<boolean> {
-  (caption as any).oldText = caption.text;
-  let text: string | undefined;
-  for (let i = 0; i < RETRY_COUNT; i += 1) {
-    try {
-      text = await translateText(
-        caption.text.replaceAll("\n", "<br>"),
-        targetLang,
-      );
-      if (text) {
-        text = text.replaceAll("<br>", "\n");
-        break;
-      }
-    } catch (error) {
-      console.warn("Re-trying caption translation", caption, error);
-    }
-  }
-  if (!text) {
-    console.error("Failed to translate caption");
-    caption.text = `(CAPTION COULD NOT BE TRANSLATED)\n${caption.text}`;
-    return false;
-  }
-  caption.text = text.trim();
-  return true;
-}
-
-async function translateCaptions(
-  captions: ContentCaption[],
-  targetLang: string,
-): Promise<boolean> {
-  // console.log("Translating", captions.length, "captions");
-  try {
-    const results: boolean[] = await Promise.all(
-      captions.map((c) => translateCaption(c, targetLang)),
-    );
-
-    const successCount = results.filter((v) => v).length;
-    const failedCount = results.length - successCount;
-    const successPercentage = (successCount / results.length) * 100;
-    const failedPercentage = (failedCount / results.length) * 100;
-    // console.log(
-    //   "Done translating captions",
-    //   results.length,
-    //   successCount,
-    //   failedCount,
-    //   successPercentage,
-    //   failedPercentage,
-    // );
-
-    if (failedPercentage > successPercentage) {
-      throw new Error("Success percentage is not acceptable");
-    }
-  } catch (error) {
-    console.error("Could not translate", captions.length, "captions", error);
-    return false;
-  }
-  return true;
-}
-
-async function translateSRTData(
-  data: string,
-  targetLang: string,
-): Promise<string | undefined> {
-  let captions: Caption[];
-  try {
-    captions = subsrt.parse(data);
-  } catch (error) {
-    console.error("Failed to parse subtitle data", error);
-    return undefined;
-  }
-
-  let translatedCaptions: Caption[] | undefined = [];
-  const contentCaptions: ContentCaption[] = [];
-  const translatedCache: Map<string, string> = new Map<string, string>();
-
-  for (const caption of captions) {
-    translatedCaptions.push(caption);
-    if (caption.type !== "caption") {
-      continue;
-    }
-    caption.text = caption.text
-      .trim()
-      .replace("\r\n", "\n")
-      .replace("\r", "\n");
-    contentCaptions.push(caption);
-  }
-
-  for (let i = 0; i < contentCaptions.length; i += 1) {
-    if (tryUseCachedCaption(contentCaptions[i], translatedCache)) {
-      continue;
-    }
-    const batch: ContentCaption[] = [contentCaptions[i]];
-
-    let j;
-    for (j = 1; j < FETCH_RATE; j += 1) {
-      if (i + j >= contentCaptions.length) {
-        break;
-      }
-      if (tryUseCachedCaption(contentCaptions[i + j], translatedCache)) {
-        continue;
-      }
-      batch.push(contentCaptions[i + j]);
-    }
-    i += j;
-
-    if (!(await translateCaptions(batch, targetLang))) {
-      translatedCaptions = undefined;
-      break;
-    }
-
-    batch.forEach((c) => translatedCache.set((c as any).oldText!, c.text));
-  }
-
-  return translatedCaptions
-    ? subsrt.build(translatedCaptions, { format: "srt" })
-    : undefined;
-}
-
-// TODO: make this support multiple providers rather than just google translate
-export async function translateSubtitle(
-  id: string,
-  srtData: string,
-  targetLang: string,
-): Promise<string | undefined> {
-  const cacheID = `${id}_${targetLang}`;
-
-  const cachedData: ArrayBuffer | undefined = SUBTITLES_CACHE.get(cacheID);
-  if (cachedData) {
-    // console.log("Using cached translation for", id, cacheID);
-    return decompressStr(cachedData);
-  }
-
-  // console.log("Translating", id);
-  const translatedData: string | undefined = await translateSRTData(
-    srtData,
-    targetLang,
-  );
-  if (!translatedData) {
-    return undefined;
-  }
-
-  // console.log("Caching translation for", id, cacheID);
-  SUBTITLES_CACHE.set(cacheID, await compressStr(translatedData));
-  return translatedData;
-}

src/utils/translation/googletranslate.ts

@@ -0,0 +1,72 @@
+import { TranslateService } from ".";
+
+const SINGLE_API_URL =
+  "https://translate.googleapis.com/translate_a/single?client=gtx&dt=t&dj=1&ie=UTF-8&oe=UTF-8&sl=auto";
+const BATCH_API_URL = "https://translate-pa.googleapis.com/v1/translateHtml";
+const BATCH_API_KEY = "AIzaSyATBXajvzQLTDHEQbcpq0Ihe0vWDHmO520";
+
+export default {
+  getName() {
+    return "Google Translate";
+  },
+
+  getConfig() {
+    return {
+      singleBatchSize: 15,
+      multiBatchSize: 80,
+      maxRetryCount: 3,
+      batchSleepMs: 200,
+    };
+  },
+
+  async translate(str, targetLang) {
+    if (!str) {
+      return "";
+    }
+
+    const response = await (
+      await fetch(
+        `${SINGLE_API_URL}&tl=${targetLang}&q=${encodeURIComponent(str)}`,
+        {
+          method: "GET",
+          headers: {
+            Accept: "application/json",
+          },
+        },
+      )
+    ).json();
+
+    if (!response.sentences) {
+      console.warn("Invalid gt response", response);
+      throw new Error("Invalid response");
+    }
+
+    return (response.sentences as any[])
+      .map((s: any) => s.trans as string)
+      .join("");
+  },
+
+  async translateMulti(batch, targetLang) {
+    if (!batch || batch.length === 0) {
+      return [];
+    }
+
+    const response = await (
+      await fetch(BATCH_API_URL, {
+        method: "POST",
+        headers: {
+          "Content-Type": "application/json+protobuf",
+          "X-goog-api-key": BATCH_API_KEY,
+        },
+        body: JSON.stringify([[batch, "auto", targetLang], "te"]),
+      })
+    ).json();
+
+    if (!Array.isArray(response) || response.length < 1) {
+      console.warn("Invalid gt batch response", response);
+      throw new Error("Invalid response");
+    }
+
+    return response[0].map((s: any) => s as string);
+  },
+} satisfies TranslateService;

src/utils/translation/index.ts

@@ -0,0 +1,187 @@
+import subsrt from "subsrt-ts";
+import { Caption, ContentCaption } from "subsrt-ts/dist/types/handler";
+
+import { Caption as PlayerCaption } from "@/stores/player/slices/source";
+
+import { compressStr, decompressStr, sleep } from "./utils";
+
+const CAPTIONS_CACHE: Map<string, ArrayBuffer> = new Map<string, ArrayBuffer>();
+
+export interface TranslateService {
+  getName(): string;
+  getConfig(): {
+    singleBatchSize: number;
+    multiBatchSize: number; // -1 = unsupported
+    maxRetryCount: number;
+    batchSleepMs: number;
+  };
+  translate(str: string, targetLang: string): Promise<string>;
+  translateMulti(batch: string[], targetLang: string): Promise<string[]>;
+}
+
+class Translator {
+  private captions: Caption[];
+
+  private contentCaptions: ContentCaption[] = [];
+
+  private contentCache: Map<string, string> = new Map<string, string>();
+
+  private targetLang: string;
+
+  private service: TranslateService;
+
+  constructor(srtData: string, targetLang: string, service: TranslateService) {
+    this.captions = subsrt.parse(srtData);
+    this.targetLang = targetLang;
+    this.service = service;
+
+    for (const caption of this.captions) {
+      if (caption.type !== "caption") {
+        continue;
+      }
+      // Normalize line endings
+      caption.text = caption.text
+        .trim()
+        .replaceAll("\r\n", "\n")
+        .replaceAll("\r", "\n");
+      this.contentCaptions.push(caption);
+    }
+  }
+
+  fillContentFromCache(content: ContentCaption): boolean {
+    const text: string | undefined = this.contentCache.get(content.text);
+    if (text) {
+      content.text = text;
+      return true;
+    }
+    return false;
+  }
+
+  async translateContent(content: ContentCaption): Promise<boolean> {
+    let result;
+    let attempts = 0;
+    const errors: any[] = [];
+
+    while (!result && attempts < 3) {
+      try {
+        result = await this.service.translate(content.text, this.targetLang);
+      } catch (err) {
+        console.warn("Translation attempt failed");
+        errors.push(err);
+        await sleep(500);
+        attempts += 1;
+      }
+    }
+
+    if (!result) {
+      console.warn("Translation failed", errors);
+      return false;
+    }
+
+    content.text = result;
+    this.contentCache.set(content.text, result);
+    return true;
+  }
+
+  async translateContentBatch(batch: ContentCaption[]): Promise<boolean> {
+    try {
+      const result = await this.service.translateMulti(
+        batch.map((content) => content.text),
+        this.targetLang,
+      );
+
+      if (result.length !== batch.length) {
+        console.warn(
+          "Batch translation size mismatch",
+          result.length,
+          batch.length,
+        );
+        return false;
+      }
+
+      for (let i = 0; i < batch.length; i += 1) {
+        batch[i].text = result[i];
+        this.contentCache.set(batch[i].text, result[i]);
+      }
+
+      return true;
+    } catch (err) {
+      console.warn("Batch translation failed", err);
+      return false;
+    }
+  }
+
+  takeBatch(): ContentCaption[] {
+    const batch: ContentCaption[] = [];
+    const batchSize =
+      this.service.getConfig().multiBatchSize === -1
+        ? this.service.getConfig().singleBatchSize
+        : this.service.getConfig().multiBatchSize;
+
+    let count = 0;
+    while (count < batchSize && this.contentCaptions.length > 0) {
+      const content: ContentCaption = this.contentCaptions.shift()!;
+      if (this.fillContentFromCache(content)) {
+        continue;
+      }
+      batch.push(content);
+      count += 1;
+    }
+
+    return batch;
+  }
+
+  async translate(): Promise<string | undefined> {
+    let batch: ContentCaption[] = this.takeBatch();
+    while (batch.length > 0) {
+      let result: boolean;
+      console.info("Translating captions batch", batch.length, batch);
+
+      if (this.service.getConfig().multiBatchSize === -1) {
+        result = (
+          await Promise.all(
+            batch.map((content) => this.translateContent(content)),
+          )
+        ).every((res) => res);
+      } else {
+        result = await this.translateContentBatch(batch);
+      }
+
+      if (!result) {
+        console.error(
+          "Failed to translate captions batch",
+          batch.length,
+          batch,
+        );
+        return undefined;
+      }
+
+      batch = this.takeBatch();
+      await sleep(this.service.getConfig().batchSleepMs);
+    }
+    return subsrt.build(this.captions, { format: "srt" });
+  }
+}
+
+export async function translate(
+  caption: PlayerCaption,
+  targetLang: string,
+  service: TranslateService,
+): Promise<string | undefined> {
+  const cacheID = `${caption.id}_${targetLang}`;
+
+  const cachedData: ArrayBuffer | undefined = CAPTIONS_CACHE.get(cacheID);
+  if (cachedData) {
+    return decompressStr(cachedData);
+  }
+
+  const translator = new Translator(caption.srtData, targetLang, service);
+
+  const result = await translator.translate();
+  if (!result) {
+    return undefined;
+  }
+
+  CAPTIONS_CACHE.set(cacheID, await compressStr(result));
+  return result;
+}

src/utils/translation/utils.ts

@@ -0,0 +1,24 @@
+export async function compressStr(string: string): Promise<ArrayBuffer> {
+  const byteArray = new TextEncoder().encode(string);
+  const cs = new CompressionStream("deflate");
+  const writer = cs.writable.getWriter();
+  writer.write(byteArray);
+  writer.close();
+  return new Response(cs.readable).arrayBuffer();
+}
+
+export async function decompressStr(byteArray: ArrayBuffer): Promise<string> {
+  const cs = new DecompressionStream("deflate");
+  const writer = cs.writable.getWriter();
+  writer.write(byteArray);
+  writer.close();
+  return new Response(cs.readable).arrayBuffer().then((arrayBuffer) => {
+    return new TextDecoder().decode(arrayBuffer);
+  });
+}
+
+export function sleep(ms: number): Promise<void> {
+  return new Promise((resolve) => {
+    setTimeout(resolve, ms);
+  });
+}