diora-web/podcasts/transcription.py

"""Transcription settings and the test of the configured Whisper instance
(docs/transcription.md). The job itself is podcasts/transcribe.py."""

import io
import json
import re
import wave
from urllib.parse import urlsplit

import requests
from django.conf import settings
from django.http import JsonResponse
from django.utils import timezone
from django.views.decorators.csrf import csrf_exempt
from django.views.decorators.http import require_http_methods

from feeds import logins

from . import chapters
from .models import EpisodeArticle, PodcastEpisode, TranscriptionJob, TranscriptionSettings

DEFAULT_MODEL = 'whisper-1'
TEST_TIMEOUT = 60            # a cold CPU instance may load its model first
LANGUAGE_RE = re.compile(r'^[a-z]{2,3}(-[A-Za-z]{2,4})?$')


def _auth(request):
    if not request.user.is_authenticated:
        return JsonResponse({'error': 'authentication required'}, status=401)
    return None


def _body(request):
    try:
        body = json.loads(request.body or b'{}')
    except (ValueError, UnicodeDecodeError):
        return None
    return body if isinstance(body, dict) else None


def for_user(user):
    """The user's settings (unsaved defaults when there are none)."""
    return TranscriptionSettings.objects.filter(user=user).first() or TranscriptionSettings(user=user)


def token_of(conf):
    """The plain API token ('' when there is none or it can't be read)."""
    return (logins.decrypt(conf.token_encrypted) or '') if conf.token_encrypted else ''


def _data(conf):
    return {
        'engine': conf.engine, 'url': conf.url, 'model': conf.model, 'language': conf.language,
        'parallel': conf.parallel, 'has_token': bool(conf.token_encrypted),
        'builtin_available': bool(settings.WHISPER_URL),
    }


@csrf_exempt
@require_http_methods(['GET', 'POST'])
def transcription_settings(request):
    if (err := _auth(request)):
        return err
    conf = for_user(request.user)
    if request.method == 'GET':
        return JsonResponse(_data(conf))
    body = _body(request)
    if body is None:
        return JsonResponse({'error': 'invalid JSON'}, status=400)
    if 'engine' in body:
        if body['engine'] not in dict(TranscriptionSettings.ENGINES):
            return JsonResponse({'error': 'engine must be builtin or custom'}, status=400)
        conf.engine = body['engine']
    if 'url' in body:
        url = body['url']
        if not isinstance(url, str):
            return JsonResponse({'error': 'url must be text'}, status=400)
        url = url.strip().rstrip('/')
        if url and (urlsplit(url).scheme not in ('http', 'https') or not urlsplit(url).hostname):
            return JsonResponse({'error': 'url must start with http:// or https://'}, status=400)
        if len(url) > 500:
            return JsonResponse({'error': 'url too long'}, status=400)
        conf.url = url
    if 'model' in body:
        if not isinstance(body['model'], str) or len(body['model'].strip()) > 100:
            return JsonResponse({'error': 'invalid model'}, status=400)
        conf.model = body['model'].strip()
    if 'language' in body:
        lang = body['language']
        if not isinstance(lang, str) or (lang.strip() and not LANGUAGE_RE.match(lang.strip())):
            return JsonResponse({'error': 'language must be empty or an ISO code'}, status=400)
        conf.language = lang.strip()
    if 'parallel' in body:
        n = body['parallel']
        if isinstance(n, bool) or not isinstance(n, int) or not 1 <= n <= 8:
            return JsonResponse({'error': 'parallel must be 1 to 8'}, status=400)
        conf.parallel = n
    if 'token' in body:
        if not isinstance(body['token'], str):
            return JsonResponse({'error': 'token must be text'}, status=400)
        conf.token_encrypted = logins.encrypt(body['token']) if body['token'] else ''
    if conf.engine == TranscriptionSettings.CUSTOM and not conf.url:
        return JsonResponse({'error': 'a custom instance needs a url'}, status=400)
    conf.save()
    return JsonResponse(_data(conf))


def silent_wav(seconds=1, rate=16000):
    """`seconds` of 16-bit mono silence as a WAV file, made in memory."""
    buf = io.BytesIO()
    with wave.open(buf, 'wb') as w:
        w.setnchannels(1)
        w.setsampwidth(2)
        w.setframerate(rate)
        w.writeframes(b'\0\0' * rate * seconds)
    return buf.getvalue()


@csrf_exempt
@require_http_methods(['POST'])
def transcription_test(request):
    """One second of silence to the configured instance: {ok, model, error}."""
    if (err := _auth(request)):
        return err
    conf = for_user(request.user)
    builtin = conf.engine == TranscriptionSettings.BUILTIN
    base = settings.WHISPER_URL if builtin else conf.url
    model = conf.model or DEFAULT_MODEL
    if not base:
        error = 'Keine eingebaute Instanz eingerichtet' if builtin else 'Keine Adresse eingetragen'
        return JsonResponse({'ok': False, 'model': model, 'error': error})
    url = base.rstrip('/') + '/audio/transcriptions'
    headers = {'User-Agent': 'diora transcription test'}
    if not builtin and conf.token_encrypted:
        headers['Authorization'] = 'Bearer ' + token_of(conf)
    files = {'file': ('silence.wav', silent_wav(), 'audio/wav')}
    data = {'model': model, 'response_format': 'json'}
    if conf.language:
        data['language'] = conf.language
    try:
        if builtin:
            # The operator's own instance (DIORA_WHISPER_URL) may sit on a private network.
            r = requests.post(url, headers=headers, files=files, data=data, timeout=TEST_TIMEOUT)
            if r.status_code >= 400:
                raise chapters.ChaptersInvalid(f'HTTP {r.status_code}')
        else:
            r = chapters._open(url, headers, method='POST', timeout=TEST_TIMEOUT, files=files, data=data)
        r.close()
    except (chapters.ChaptersInvalid, chapters.ChaptersUnavailable) as e:
        return JsonResponse({'ok': False, 'model': model, 'error': str(e)})
    except requests.RequestException as e:
        return JsonResponse({'ok': False, 'model': model, 'error': e.__class__.__name__})
    return JsonResponse({'ok': True, 'model': model, 'error': ''})


# ---------------------------------------------------------------------------
# Jobs
# ---------------------------------------------------------------------------

CLAIM_MINUTES = 15
ACTIVE = (TranscriptionJob.QUEUED, TranscriptionJob.DOWNLOADING, TranscriptionJob.TRANSCRIBING)


def job_state(job):
    """The job as the API shows it; state `none` when there is no job."""
    if job is None:
        return {'state': TranscriptionJob.NONE, 'chunks_done': 0, 'chunks_total': 0, 'error': '',
                'want_article': False, 'started_at': None, 'finished_at': None}
    return {
        'state': job.state, 'chunks_done': job.chunks_done, 'chunks_total': job.chunks_total,
        'error': job.error, 'want_article': job.want_article,
        'started_at': job.started_at.isoformat() if job.started_at else None,
        'finished_at': job.finished_at.isoformat() if job.finished_at else None,
    }


def job_of(episode):
    return TranscriptionJob.objects.filter(episode=episode).first()


def _too_large(url):
    """Whether the server says the file at `url` is over the size limit (a failed look doesn't count)."""
    from . import silence
    try:
        r = requests.head(url, timeout=5, allow_redirects=True, headers={'User-Agent': 'diora transcription'})
        return int(r.headers.get('Content-Length') or 0) > silence.MAX_BYTES
    except (requests.RequestException, ValueError):
        return False


@csrf_exempt
@require_http_methods(['GET', 'POST', 'DELETE'])
def episode_transcribe(request, pk):
    """POST starts the episode's transcription (or just raises `want_article` on one that exists),
    GET tells how far it is, DELETE cancels a queued or running one."""
    from . import transcribe, views
    from .youtube import is_video
    if (err := _auth(request)):
        return err
    ep = PodcastEpisode.objects.select_related('feed').filter(pk=pk, feed__user=request.user).first()
    if ep is None:
        return JsonResponse({'error': 'not found'}, status=404)
    job = job_of(ep)

    if request.method == 'GET':
        transcribe.kick()   # after a restart the first request brings the work back
        return JsonResponse(job_state(job))

    if request.method == 'DELETE':
        if job is not None and job.state != TranscriptionJob.DONE:
            job.delete()    # a running worker notices its row is gone and stops
            job = None
        return JsonResponse(job_state(job))

    body = _body(request)
    if body is None:
        return JsonResponse({'error': 'invalid JSON'}, status=400)
    article = body.get('article') is True
    if is_video(ep.audio_url):
        return JsonResponse({'error': 'youtube'}, status=400)

    has_feed = views._has_transcript(ep)
    if job is not None and (job.state in ACTIVE or (job.state == TranscriptionJob.DONE and (job.data or has_feed))):
        if article and not job.want_article:
            job.want_article = True
            job.save(update_fields=['want_article'])
        return JsonResponse(job_state(job))
    if has_feed:   # the feed's own transcript is there: only the article is wanted
        if job is None:
            job = TranscriptionJob(episode=ep)
        job.state, job.want_article, job.finished_at = TranscriptionJob.DONE, article or job.want_article, timezone.now()
        job.save()
        return JsonResponse(job_state(job))

    if not ep.audio_url:
        return JsonResponse({'error': 'no audio'}, status=400)
    if _too_large(ep.audio_url):
        return JsonResponse({'error': 'too large'}, status=400)
    if job is not None:
        job.delete()   # a failed or hollow earlier run: the new one starts at the back of the queue
    job = TranscriptionJob.objects.create(episode=ep, state=TranscriptionJob.QUEUED, want_article=article)
    transcribe.kick()
    return JsonResponse(job_state(job))


def _wishes(user):
    """The user's jobs whose transcript is there and whose article is wanted but not written."""
    return (TranscriptionJob.objects
            .filter(episode__feed__user=user, state=TranscriptionJob.DONE, want_article=True)
            .exclude(episode__in=EpisodeArticle.objects.values('episode')))


@csrf_exempt
@require_http_methods(['GET'])
def transcription_pending(request):
    """The open wishes for an article that no device holds right now."""
    from . import transcribe
    if (err := _auth(request)):
        return err
    transcribe.kick()
    jobs = (_wishes(request.user).exclude(claimed_until__gt=timezone.now())
            .select_related('episode__feed').order_by('id'))
    return JsonResponse([{
        'episode_id': j.episode_id, 'title': j.episode.title, 'podcast': j.episode.feed.title,
        'duration_seconds': j.episode.duration_seconds,
    } for j in jobs], safe=False)


@csrf_exempt
@require_http_methods(['POST', 'DELETE'])
def article_claim(request, pk):
    """Reserve the wish for an article for the calling device for CLAIM_MINUTES (POST), or give
    it up (DELETE). The device names itself with `device` in the JSON body or the
    X-Diora-Device header. 409 when another device holds it."""
    from datetime import timedelta
    if (err := _auth(request)):
        return err
    body = _body(request)
    if body is None:
        return JsonResponse({'error': 'invalid JSON'}, status=400)
    device = body.get('device') or request.META.get('HTTP_X_DIORA_DEVICE') or ''
    if not isinstance(device, str) or not device.strip():
        return JsonResponse({'error': 'device missing'}, status=400)
    device = device.strip()[:60]
    job = _wishes(request.user).filter(episode_id=pk).first()
    if job is None:
        return JsonResponse({'error': 'no open wish'}, status=404)
    now = timezone.now()
    if job.claimed_until and job.claimed_until > now and job.claimed_by != device:
        return JsonResponse({'error': 'claimed', 'claimed_until': job.claimed_until.isoformat()}, status=409)
    if request.method == 'DELETE':
        job.claimed_by, job.claimed_until = '', None
    else:
        job.claimed_by, job.claimed_until = device, now + timedelta(minutes=CLAIM_MINUTES)
    job.save(update_fields=['claimed_by', 'claimed_until'])
    return JsonResponse({'claimed_by': job.claimed_by,
                         'claimed_until': job.claimed_until.isoformat() if job.claimed_until else None})