feat: translate subtitles with an OpenAI-compatible LLM

src/translate/translationservice.cpp

@@ -0,0 +1,477 @@
+#include "translationservice.h"
+
+#include "apikeystore.h"
+#include "qtpotatosettings.h"
+
+#include <KLocalizedString>
+
+#include <QCoroFuture>
+#include <QCoroTimer>
+
+#include <QCryptographicHash>
+#include <QDir>
+#include <QFileInfo>
+#include <QHash>
+#include <QLocale>
+#include <QNetworkAccessManager>
+#include <QPointer>
+#include <QSaveFile>
+#include <QStandardPaths>
+#include <QtConcurrentRun>
+
+#include <chrono>
+#include <vector>
+
+using namespace Qt::Literals::StringLiterals;
+using namespace std::chrono_literals;
+
+namespace
+{
+// Source lines before a batch that the model sees for context.
+constexpr qsizetype ContextLines = 6;
+constexpr int MaxAttempts = 3;
+// Requests in flight at once. Batches only use source lines as context, so they don't
+// depend on each other.
+constexpr int ParallelRequests = 3;
+// Bump when the prompt changes so old cached translations are not reused.
+constexpr auto CacheVersion = "llm-v1";
+
+// "Vietnamese" -> "vi", so the file ends in ".vi.srt" like other players expect.
+QString languageSuffix(const QString &language)
+{
+    for (int i = QLocale::C + 1; i <= QLocale::LastLanguage; ++i) {
+        const auto candidate = static_cast<QLocale::Language>(i);
+        if (QLocale::languageToString(candidate).compare(language, Qt::CaseInsensitive) == 0) {
+            const QString code = QLocale::languageToCode(candidate, QLocale::ISO639Part1);
+            return code.isEmpty() ? QLocale::languageToCode(candidate) : code;
+        }
+    }
+    QString suffix = language.toLower();
+    suffix.removeIf([](QChar c) {
+        return !c.isLetterOrNumber() && c != u'-';
+    });
+    return suffix.isEmpty() ? u"translated"_s : suffix;
+}
+
+QString cachePath(const QUrl &videoUrl, const QString &key, const QString &language)
+{
+    const QString dir = QStandardPaths::writableLocation(QStandardPaths::CacheLocation) + u"/translations"_s;
+    const QString baseName = videoUrl.isLocalFile() ? QFileInfo(videoUrl.toLocalFile()).completeBaseName() : u"subtitles"_s;
+    return u"%1/%2.%3.%4.srt"_s.arg(dir, baseName, key, languageSuffix(language));
+}
+}
+
+struct TranslationService::Job {
+    std::atomic_bool cancelled = false;
+    std::unique_ptr<LlmClient> client;
+    QString mediaTitle;
+};
+
+struct TranslationService::Batches {
+    QStringList lines;
+    QStringList translations;
+    qsizetype batchSize = 0;
+    qsizetype nextStart = 0;
+    qsizetype doneLines = 0;
+    QString error;
+};
+
+TranslationService::TranslationService(QObject *parent)
+    : QObject(parent)
+    , m_network(new QNetworkAccessManager(this))
+{
+}
+
+TranslationService::~TranslationService()
+{
+    cancel();
+}
+
+bool TranslationService::isBusy() const
+{
+    return m_busy;
+}
+
+double TranslationService::progress() const
+{
+    return m_progress;
+}
+
+QString TranslationService::statusText() const
+{
+    return m_statusText;
+}
+
+QString TranslationService::errorText() const
+{
+    return m_errorText;
+}
+
+QStringList TranslationService::availableModels() const
+{
+    return m_availableModels;
+}
+
+bool TranslationService::isFetchingModels() const
+{
+    return m_fetchingModels;
+}
+
+void TranslationService::translate(const QUrl &videoUrl, const QVariantMap &track, const QString &mediaTitle)
+{
+    cancel();
+    m_job = std::make_shared<Job>();
+    m_job->mediaTitle = mediaTitle;
+    run(videoUrl, track);
+}
+
+void TranslationService::cancel()
+{
+    if (!m_job) {
+        return;
+    }
+    // The running coroutine keeps its own reference and stops at its next check.
+    m_job->cancelled = true;
+    if (m_job->client) {
+        m_job->client->abort();
+    }
+    m_job.reset();
+    setBusy(false);
+    setStatusText({});
+}
+
+void TranslationService::fetchModels()
+{
+    if (!m_fetchingModels) {
+        runFetchModels();
+    }
+}
+
+void TranslationService::clearError()
+{
+    setErrorText({});
+}
+
+QCoro::Task<std::pair<LlmClient::Config, QString>> TranslationService::loadConfig()
+{
+    const auto settings = QtPotatoSettings::self();
+    LlmClient::Config config{
+        LlmClient::normalizedBaseUrl(settings->apiBaseUrl()),
+        {},
+        settings->model().trimmed(),
+        settings->targetLanguage().trimmed(),
+    };
+
+    if (!config.baseUrl.isValid() || config.baseUrl.scheme().isEmpty() || config.baseUrl.host().isEmpty()) {
+        co_return std::pair{config, i18n("Set the API base URL in Settings first.")};
+    }
+
+    const ApiKeyStore::ReadResult key = co_await ApiKeyStore::readKey();
+    if (!key.error.isEmpty()) {
+        co_return std::pair{config, i18n("Could not read the API key from the wallet: %1", key.error)};
+    }
+    config.apiKey = key.key;
+    co_return std::pair{config, QString()};
+}
+
+QCoro::Task<Subtitles::ExtractResult> TranslationService::loadCues(std::shared_ptr<Job> job, QUrl videoUrl, QVariantMap track)
+{
+    if (track.value(u"external"_s).toBool()) {
+        const QString path = track.value(u"externalFilename"_s).toString();
+        co_return co_await QtConcurrent::run([path] {
+            QString error;
+            QList<SubtitleCue> cues = Subtitles::parseFile(path, &error);
+            return Subtitles::ExtractResult{cues, error};
+        });
+    }
+
+    if (!videoUrl.isLocalFile()) {
+        co_return Subtitles::ExtractResult{{}, i18n("Only subtitles of local files can be translated.")};
+    }
+    const int streamIndex = track.value(u"ffIndex"_s, -1).toInt();
+    if (streamIndex < 0) {
+        co_return Subtitles::ExtractResult{{}, i18n("mpv did not report where this subtitle track is stored.")};
+    }
+
+    const QString path = videoUrl.toLocalFile();
+    co_return co_await QtConcurrent::run([path, streamIndex, job] {
+        return Subtitles::extractEmbedded(path, streamIndex, job->cancelled);
+    });
+}
+
+QCoro::Task<> TranslationService::run(QUrl videoUrl, QVariantMap track)
+{
+    const std::shared_ptr<Job> job = m_job;
+
+    setErrorText({});
+    setProgress(0);
+    setBusy(true);
+    setStatusText(i18n("Reading subtitles…"));
+
+    const auto [config, configError] = co_await loadConfig();
+    if (job->cancelled) {
+        co_return;
+    }
+    if (!configError.isEmpty()) {
+        fail(job, configError);
+        co_return;
+    }
+    if (config.model.isEmpty()) {
+        fail(job, i18n("Choose a model in Settings first."));
+        co_return;
+    }
+
+    Subtitles::ExtractResult source = co_await loadCues(job, videoUrl, track);
+    if (job->cancelled) {
+        co_return;
+    }
+    if (!source.error.isEmpty()) {
+        fail(job, source.error);
+        co_return;
+    }
+    const QList<SubtitleCue> cues = Subtitles::mergeRepeated(source.cues);
+    if (cues.isEmpty()) {
+        fail(job, i18n("This subtitle track has no text to translate."));
+        co_return;
+    }
+
+    // Same subtitles, model and language give the same file, so it is translated once.
+    QCryptographicHash hash(QCryptographicHash::Sha1);
+    hash.addData(CacheVersion);
+    hash.addData(config.model.toUtf8());
+    hash.addData(config.targetLanguage.toUtf8());
+    hash.addData(Subtitles::toSrt(cues).toUtf8());
+    const QString outputPath = cachePath(videoUrl, QString::fromLatin1(hash.result().toHex().left(12)), config.targetLanguage);
+
+    if (QFileInfo::exists(outputPath)) {
+        m_job.reset();
+        setBusy(false);
+        setStatusText({});
+        Q_EMIT translated(QUrl::fromLocalFile(outputPath), true);
+        co_return;
+    }
+
+    // Repeated lines (songs, catchphrases) are sent once.
+    QStringList uniqueLines;
+    QHash<QString, qsizetype> uniqueIndex;
+    QList<qsizetype> cueToUnique;
+    for (const SubtitleCue &cue : cues) {
+        auto it = uniqueIndex.constFind(cue.text);
+        if (it == uniqueIndex.cend()) {
+            it = uniqueIndex.insert(cue.text, uniqueLines.size());
+            uniqueLines.append(cue.text);
+        }
+        cueToUnique.append(*it);
+    }
+
+    job->client = std::make_unique<LlmClient>(m_network, config);
+    auto batches = std::make_shared<Batches>();
+    batches->lines = uniqueLines;
+    batches->translations.resize(uniqueLines.size());
+    batches->batchSize = QtPotatoSettings::self()->batchSize();
+    setStatusText(i18n("Translating… %1 of %2 lines", 0, uniqueLines.size()));
+
+    // Tasks start running as soon as they are created; awaiting them just waits for the end.
+    std::vector<QCoro::Task<>> workers;
+    for (int i = 0; i < ParallelRequests; ++i) {
+        workers.push_back(translateBatches(job, batches));
+    }
+    for (QCoro::Task<> &worker : workers) {
+        co_await std::move(worker);
+    }
+    if (job->cancelled) {
+        co_return;
+    }
+    if (!batches->error.isEmpty()) {
+        fail(job, batches->error);
+        co_return;
+    }
+
+    QList<SubtitleCue> output = cues;
+    for (qsizetype i = 0; i < output.size(); ++i) {
+        output[i].text = batches->translations.at(cueToUnique.at(i));
+    }
+
+    QDir().mkpath(QFileInfo(outputPath).absolutePath());
+    QSaveFile file(outputPath);
+    if (!file.open(QIODevice::WriteOnly) || file.write(Subtitles::toSrt(output).toUtf8()) < 0 || !file.commit()) {
+        fail(job, i18n("Could not save the translation to %1: %2", outputPath, file.errorString()));
+        co_return;
+    }
+
+    m_job.reset();
+    setBusy(false);
+    setStatusText({});
+    Q_EMIT translated(QUrl::fromLocalFile(outputPath), false);
+}
+
+QCoro::Task<LlmClient::BatchResult> TranslationService::translateLines(std::shared_ptr<Job> job, QStringList lines, QStringList context)
+{
+    LlmClient::BatchResult result;
+    for (int attempt = 1; attempt <= MaxAttempts; ++attempt) {
+        result = co_await job->client->translate(lines, context, job->mediaTitle);
+        if (job->cancelled || result.error.isEmpty() || !result.retryable || attempt == MaxAttempts) {
+            break;
+        }
+        co_await QCoro::sleepFor(2s * attempt);
+        if (job->cancelled) {
+            co_return result;
+        }
+    }
+    if (job->cancelled || !result.error.isEmpty()) {
+        co_return result;
+    }
+
+    QList<qsizetype> missing;
+    for (qsizetype i = 0; i < result.translations.size(); ++i) {
+        if (result.translations.at(i).isNull()) {
+            missing.append(i);
+        }
+    }
+    if (missing.isEmpty()) {
+        co_return result;
+    }
+    if (lines.size() == 1) {
+        co_return LlmClient::BatchResult{{}, i18n("The model did not translate this line: %1", lines.first()), false};
+    }
+
+    if (missing.size() == lines.size()) {
+        // Nothing came back usable; smaller requests are easier for the model.
+        const qsizetype half = lines.size() / 2;
+        const LlmClient::BatchResult first = co_await translateLines(job, lines.first(half), context);
+        if (job->cancelled || !first.error.isEmpty()) {
+            co_return first;
+        }
+        QStringList secondContext = context + lines.first(half);
+        secondContext = secondContext.mid(std::max<qsizetype>(0, secondContext.size() - ContextLines));
+        const LlmClient::BatchResult second = co_await translateLines(job, lines.sliced(half), secondContext);
+        if (job->cancelled || !second.error.isEmpty()) {
+            co_return second;
+        }
+        co_return LlmClient::BatchResult{first.translations + second.translations, {}, false};
+    }
+
+    // Ask again only for the lines the model skipped.
+    QStringList retryLines;
+    for (const qsizetype i : missing) {
+        retryLines.append(lines.at(i));
+    }
+    const LlmClient::BatchResult retry = co_await translateLines(job, retryLines, context);
+    if (job->cancelled || !retry.error.isEmpty()) {
+        co_return retry;
+    }
+    for (qsizetype i = 0; i < missing.size(); ++i) {
+        result.translations[missing.at(i)] = retry.translations.at(i);
+    }
+    co_return result;
+}
+
+QCoro::Task<> TranslationService::translateBatches(std::shared_ptr<Job> job, std::shared_ptr<Batches> batches)
+{
+    // After an error, requests already sent may finish, but no new ones start.
+    while (!job->cancelled && batches->error.isEmpty() && batches->nextStart < batches->lines.size()) {
+        const qsizetype start = batches->nextStart;
+        batches->nextStart += batches->batchSize;
+
+        const qsizetype contextStart = std::max<qsizetype>(0, start - ContextLines);
+        const LlmClient::BatchResult result =
+            co_await translateLines(job, batches->lines.mid(start, batches->batchSize), batches->lines.mid(contextStart, start - contextStart));
+        if (job->cancelled) {
+            co_return;
+        }
+        if (!result.error.isEmpty()) {
+            if (batches->error.isEmpty()) {
+                batches->error = result.error;
+            }
+            co_return;
+        }
+
+        for (qsizetype i = 0; i < result.translations.size(); ++i) {
+            batches->translations[start + i] = result.translations.at(i);
+        }
+        batches->doneLines += result.translations.size();
+        setProgress(double(batches->doneLines) / batches->lines.size());
+        setStatusText(i18n("Translating… %1 of %2 lines", batches->doneLines, batches->lines.size()));
+    }
+}
+
+QCoro::Task<> TranslationService::runFetchModels()
+{
+    QPointer<TranslationService> self(this);
+    setFetchingModels(true);
+    setErrorText({});
+
+    const auto [config, configError] = co_await loadConfig();
+    if (!self) {
+        co_return;
+    }
+    if (!configError.isEmpty()) {
+        setErrorText(configError);
+        setFetchingModels(false);
+        co_return;
+    }
+
+    LlmClient client(m_network, config);
+    const LlmClient::ModelsResult result = co_await client.listModels();
+    if (!self) {
+        co_return;
+    }
+    if (!result.error.isEmpty()) {
+        setErrorText(i18n("Could not list models: %1", result.error));
+    } else if (result.models != m_availableModels) {
+        m_availableModels = result.models;
+        Q_EMIT availableModelsChanged();
+    }
+    setFetchingModels(false);
+}
+
+void TranslationService::fail(const std::shared_ptr<Job> &job, const QString &error)
+{
+    if (m_job != job) {
+        return;
+    }
+    m_job.reset();
+    setErrorText(error);
+    setBusy(false);
+    setStatusText({});
+}
+
+void TranslationService::setBusy(bool busy)
+{
+    if (busy != m_busy) {
+        m_busy = busy;
+        Q_EMIT busyChanged();
+    }
+}
+
+void TranslationService::setProgress(double progress)
+{
+    if (progress != m_progress) {
+        m_progress = progress;
+        Q_EMIT progressChanged();
+    }
+}
+
+void TranslationService::setStatusText(const QString &text)
+{
+    if (text != m_statusText) {
+        m_statusText = text;
+        Q_EMIT statusTextChanged();
+    }
+}
+
+void TranslationService::setErrorText(const QString &text)
+{
+    if (text != m_errorText) {
+        m_errorText = text;
+        Q_EMIT errorTextChanged();
+    }
+}
+
+void TranslationService::setFetchingModels(bool fetching)
+{
+    if (fetching != m_fetchingModels) {
+        m_fetchingModels = fetching;
+        Q_EMIT fetchingModelsChanged();
+    }
+}