from __future__ import annotations import asyncio import json import logging import os import signal import time import uuid import wave from dataclasses import dataclass, field from pathlib import Path from tempfile import NamedTemporaryFile from typing import Any import aiohttp import nats import numpy as np from nats.aio.msg import Msg from src.diarization import ( attribute_words, attributed_words_payload, build_speaker_turns, diarization_backend_ready, diarization_segments_payload, diarize_http_json, parse_diarization_result, render_speaker_text, speaker_turns_payload, ) from src.diarization_syntax import apply_syntax_postprocess from src.webui.file_mode import ( FileJobStore, FileModeSettings, MissingWordTimestampsError, TimedTranscriptBuilder, estimate_chunk_total, iter_normalized_wav_chunk_files, load_file_mode_settings, prepare_diarization_audio_artifact, probe_wav_duration_s, resolve_optional_bool, transcribe_http_json, transcribe_http_text, write_normalized_audio_wav, ) from src.webui.job_lock import FileJobLockClient, job_lock_enabled, page_label from src.utils.audio_utils import resample_audio from src.utils.silero_onnx_vad import SileroOnnxVAD, download_model_file, get_speech_timestamps from src.utils.wire_packet import parse_wire_packet logger = logging.getLogger("stt_whisper_to_nats") FILE_CHUNK_PROGRESS_INTERVAL_S = 5.0 FILE_JOB_TERMINAL_CLEANUP_DELAY_S = 30.0 DEFAULT_VAD_MODEL_URL = ( "https://github.com/snakers4/silero-vad/raw/refs/heads/master/src/silero_vad/data/silero_vad.onnx" ) def _env(name: str, default: str = "") -> str: return str(os.getenv(name, default) or "").strip() def _env_bool(name: str, default: bool = False) -> bool: raw = _env(name, "1" if default else "0").lower() return raw in {"1", "true", "yes", "on"} @dataclass(slots=True) class PhraseStreamState: out_subject: str session_id: str | None sample_rate: int buffer: list[np.ndarray] buffer_samples: int check_in_progress: bool = False flush_requested: bool = False lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False) class SttWhisperToNatsService: def __init__(self, settings: FileModeSettings) -> None: self.service_name = _env("STACK_SERVICE_NAME", "stt_whisper_to_nats") self.nats_url = _env("NATS_URL") self.in_prefix = _env("ASR_IN_PREFIX") self.out_prefix = _env("ASR_OUT_PREFIX") self._settings = settings self.min_phrase_ms = int(_env("MIN_PHRASE_MS", "100")) self.vad_models_dir = Path(_env("VAD_MODELS", "/app/models/vad")) self.vad_model_url = _env("VAD_MODEL_URL", DEFAULT_VAD_MODEL_URL) self.vad_sample_rate = int(_env("VAD_SAMPLE_RATE", "16000")) self.min_speech_duration_ms = int(_env("MIN_SPEECH_DURATION_MS", "250")) self.min_silence_duration_ms = int(_env("MIN_SILENCE_DURATION_MS", "500")) self.max_speech_duration_s = float(_env("MAX_SPEECH_DURATION_S", "30.0")) self.speech_pad_ms = int(_env("SPEECH_PAD_MS", "30")) self.vad_threshold = float(_env("VAD_THRESHOLD", "0.8")) self.buffer_check_interval_s = float(_env("BUFFER_CHECK_INTERVAL_S", "1.0")) self.max_concurrency = max(1, int(_env("WHISPER_MAX_CONCURRENCY", "1"))) if not self.nats_url: raise ValueError("NATS_URL is required") if not self.in_prefix: raise ValueError("ASR_IN_PREFIX is required") if not self.out_prefix: raise ValueError("ASR_OUT_PREFIX is required") if not self.in_prefix.endswith("."): raise ValueError("ASR_IN_PREFIX must end with '.'") if not self.out_prefix.endswith("."): self.out_prefix = f"{self.out_prefix}." if not self._settings.linto_http_url: raise ValueError("LINTO_HTTP_URL is required for live phrase transcription") self.in_subscribe = f"{self.in_prefix}>" self.nc: nats.NATS | None = None self.http: aiohttp.ClientSession | None = None self.stop_event = asyncio.Event() self.tasks: set[asyncio.Task] = set() self.streams: dict[str, PhraseStreamState] = {} self.streams_lock = asyncio.Lock() self._vad_model: SileroOnnxVAD | None = None self._vad_model_lock = asyncio.Lock() self._vad_run_lock = asyncio.Lock() self._semaphore = asyncio.Semaphore(self.max_concurrency) async def run(self) -> None: timeout = aiohttp.ClientTimeout(total=None, connect=10.0) self.http = aiohttp.ClientSession(timeout=timeout) self.nc = await nats.connect( servers=[self.nats_url], name=self.service_name, max_reconnect_attempts=-1, reconnect_time_wait=2, ping_interval=10, ) await self.nc.subscribe(self.in_subscribe, cb=self._handle_message) logger.info("Connected to NATS %s", self.nats_url) logger.info("Subscribed to %s", self.in_subscribe) logger.info("LinTO HTTP target %s", self._settings.linto_http_url) if self._settings.linto_http_host: logger.info("LinTO HTTP host override %s", self._settings.linto_http_host) await self.stop_event.wait() await self._shutdown() def stop(self) -> None: self.stop_event.set() async def _shutdown(self) -> None: for task in list(self.tasks): task.cancel() async with self.streams_lock: self.streams.clear() if self.nc: try: await self.nc.drain() except Exception: await self.nc.close() if self.http: await self.http.close() async def _handle_message(self, msg: Msg) -> None: task = asyncio.create_task(self._process_message(msg)) self.tasks.add(task) task.add_done_callback(self.tasks.discard) def _derive_out_subject(self, in_subject: str) -> str | None: if not in_subject.startswith(self.in_prefix): logger.error( "Dropping message: subject does not match ASR_IN_PREFIX (subject=%s prefix=%s)", in_subject, self.in_prefix, ) return None suffix = in_subject[len(self.in_prefix):] if not suffix: logger.error("Dropping message: empty token suffix in subject=%s", in_subject) return None return f"{self.out_prefix}{suffix}" async def _process_message(self, msg: Msg) -> None: out_subject = self._derive_out_subject(str(msg.subject or "")) if not out_subject: return try: meta, audio = parse_wire_packet(msg.data) except Exception as exc: logger.warning("Invalid packet on %s: %s", msg.subject, exc) await self._publish_json(out_subject, {"error": "invalid_packet", "details": str(exc)}) return msg_type = str(meta.get("type") or "").strip() if msg_type not in {"frame", "end"}: await self._publish_json( out_subject, {"error": "unsupported_packet_type", "details": f"type={msg_type or '-'}"}, ) return if msg_type == "frame": await self._handle_frame(out_subject, meta, audio) return await self._handle_end(out_subject, meta) async def _handle_frame(self, out_subject: str, meta: dict[str, Any], audio: np.ndarray) -> None: session_id = _clean_optional_text(meta.get("session_id")) sample_rate = int(meta.get("sample_rate") or self.vad_sample_rate) state = self.streams.get(out_subject) if not state: async with self.streams_lock: state = self.streams.get(out_subject) if not state: state = PhraseStreamState( out_subject=out_subject, session_id=session_id, sample_rate=sample_rate, buffer=[], buffer_samples=0, ) self.streams[out_subject] = state logger.info("Phrase stream opened: %s sample_rate=%s", out_subject, sample_rate) schedule_check = False async with state.lock: state.session_id = state.session_id or session_id state.sample_rate = sample_rate or state.sample_rate or self.vad_sample_rate if audio.size: state.buffer.append(audio) state.buffer_samples += int(audio.size) duration_s = ( state.buffer_samples / max(1, state.sample_rate) if state.buffer_samples else 0.0 ) schedule_check = duration_s >= self.buffer_check_interval_s and not state.check_in_progress if schedule_check: task = asyncio.create_task(self._check_stream_buffer(out_subject, flush=False)) self.tasks.add(task) task.add_done_callback(self.tasks.discard) async def _handle_end(self, out_subject: str, meta: dict[str, Any]) -> None: state = self.streams.get(out_subject) if not state: return async with state.lock: if not state.session_id: state.session_id = _clean_optional_text(meta.get("session_id")) state.flush_requested = True schedule_check = not state.check_in_progress if schedule_check: task = asyncio.create_task(self._check_stream_buffer(out_subject, flush=True)) self.tasks.add(task) task.add_done_callback(self.tasks.discard) async def _check_stream_buffer(self, out_subject: str, *, flush: bool) -> None: state = self.streams.get(out_subject) if not state: return async with state.lock: if state.check_in_progress: return state.check_in_progress = True effective_flush = bool(flush or state.flush_requested) snapshot_chunks = list(state.buffer) snapshot_samples = int(state.buffer_samples) snapshot_sr = int(state.sample_rate or self.vad_sample_rate) session_id = state.session_id try: if not snapshot_chunks: return audio = np.concatenate(snapshot_chunks) if snapshot_sr != self.vad_sample_rate: audio = resample_audio(audio, snapshot_sr, self.vad_sample_rate) if not effective_flush and (len(audio) / self.vad_sample_rate) < 1.0: return timestamps = await self._get_timestamps(audio) if not timestamps: if (len(audio) / self.vad_sample_rate) > 5.0: async with state.lock: del state.buffer[: len(snapshot_chunks)] state.buffer_samples = max(0, int(state.buffer_samples) - snapshot_samples) return should_flush = effective_flush if not should_flush: last_end_sample = int(timestamps[-1]["end"]) silence_after = (len(audio) - last_end_sample) / self.vad_sample_rate should_flush = silence_after >= (self.min_silence_duration_ms / 1000.0) if not should_flush: return for ts in timestamps: start_sample = int(ts["start"]) end_sample = int(ts["end"]) segment_audio = audio[start_sample:end_sample] if not segment_audio.size: continue phrase_id = str(uuid.uuid4()) task = asyncio.create_task( self._transcribe_and_publish( reply_subject=out_subject, session_id=session_id, phrase_id=phrase_id, audio=segment_audio, sample_rate=self.vad_sample_rate, ) ) self.tasks.add(task) task.add_done_callback(self.tasks.discard) async with state.lock: del state.buffer[: len(snapshot_chunks)] state.buffer_samples = max(0, int(state.buffer_samples) - snapshot_samples) finally: needs_final_flush = False async with state.lock: state.check_in_progress = False needs_final_flush = bool(state.flush_requested and not flush) if effective_flush: async with self.streams_lock: self.streams.pop(out_subject, None) logger.info("Phrase stream closed: %s", out_subject) elif needs_final_flush: task = asyncio.create_task(self._check_stream_buffer(out_subject, flush=True)) self.tasks.add(task) task.add_done_callback(self.tasks.discard) async def _transcribe_and_publish( self, *, reply_subject: str, session_id: str | None, phrase_id: str, audio: np.ndarray, sample_rate: int, ) -> None: if not self.http: raise RuntimeError("HTTP session is not initialized") duration_s = float(len(audio) / max(1, sample_rate)) min_samples = int(sample_rate * self.min_phrase_ms / 1000) if audio.size < min_samples: audio = np.pad(audio, (0, min_samples - audio.size), mode="constant") pending_payload: dict[str, Any] = {"pending": True, "phrase_id": phrase_id} pending_payload.update(self._session_field(session_id)) await self._publish_json(reply_subject, pending_payload) wav_path: Path | None = None try: wav_path = await asyncio.to_thread(self._write_phrase_wav, phrase_id, audio, sample_rate) async with self._semaphore: text = await transcribe_http_text( settings=self._settings, file_path=wav_path, file_name=f"{phrase_id}.wav", language=self._settings.linto_http_language or None, http=self.http, ) except Exception as exc: logger.exception("Phrase transcription failed for %s: %s", reply_subject, exc) await self._publish_json( reply_subject, { "phrase_id": phrase_id, "error": "transcription_failed", "details": str(exc), **self._session_field(session_id), }, ) return finally: if wav_path is not None: try: wav_path.unlink(missing_ok=True) except Exception: pass cleaned_text = str(text or "").strip() if not cleaned_text: logger.info("Ignoring empty phrase result for %s phrase_id=%s", reply_subject, phrase_id) return await self._publish_json( reply_subject, { "phrase_id": phrase_id, "text": cleaned_text, "duration": duration_s, **self._session_field(session_id), }, ) async def _get_timestamps(self, audio_array: np.ndarray) -> list[dict[str, Any]]: await self._ensure_vad_loaded() if not self._vad_model: return [] async with self._vad_run_lock: return await asyncio.to_thread( get_speech_timestamps, audio_array, self._vad_model, threshold=self.vad_threshold, sampling_rate=self.vad_sample_rate, min_speech_duration_ms=self.min_speech_duration_ms, min_silence_duration_ms=self.min_silence_duration_ms, max_speech_duration_s=self.max_speech_duration_s, speech_pad_ms=self.speech_pad_ms, return_seconds=False, ) async def _ensure_vad_loaded(self) -> None: if self._vad_model: return async with self._vad_model_lock: if self._vad_model: return filename = os.path.basename(self.vad_model_url) base_path = self.vad_models_dir.expanduser() model_path = base_path if base_path.suffix.lower() == ".onnx" else (base_path / filename) model_path.parent.mkdir(parents=True, exist_ok=True) if not model_path.exists(): await asyncio.to_thread(download_model_file, str(model_path), self.vad_model_url) self._vad_model = await asyncio.to_thread(SileroOnnxVAD, str(model_path)) logger.info("Silero VAD model loaded from %s", model_path) @staticmethod def _write_phrase_wav(phrase_id: str, audio: np.ndarray, sample_rate: int) -> Path: pcm = np.clip(audio, -1.0, 1.0) pcm_i16 = (pcm * 32767.0).astype(" dict[str, str]: if not session_id: return {} return {"session_id": session_id} async def _publish_json(self, subject: str, payload: dict[str, Any]) -> None: if not self.nc: return await self.nc.publish(subject, json.dumps(payload, ensure_ascii=False).encode("utf-8")) def _clean_optional_text(value: object) -> str | None: if not isinstance(value, str): return None cleaned = value.strip() return cleaned or None def _normalize_diarization_runtime_label(value: object) -> str: raw = str(value or "").strip() if not raw: return "" lowered = raw.lower() if lowered in {"gpu", "cuda"}: return "GPU" if lowered == "cpu": return "CPU" return raw.upper() def _diarization_wait_details(runtime_label: object) -> str: label = _normalize_diarization_runtime_label(runtime_label) if label: return ( f"Plain text ready. Speaker separation is still running on {label} " "and can take several minutes for longer files." ) return "Plain text ready. Speaker separation is still running and can take several minutes for longer files." def _diar_text_done_payload( *, job_id: str, file_name: str, duration_s: float, speaker_count: int, speaker_text: str, speaker_turns: list[dict[str, Any]], speaker_text_raw: str, speaker_turns_raw: list[dict[str, Any]], speaker_text_syntax: str, speaker_turns_syntax: list[dict[str, Any]], syntax_postprocess_enabled: bool, syntax_postprocess_applied: bool, syntax_sentence_count: int, syntax_low_confidence_count: int, session_id: str | None, ) -> dict[str, Any]: payload: dict[str, Any] = { "type": "diar_text_done", "job_id": job_id, "file_name": file_name, "phase": "done", "duration_s": round(duration_s, 3), "speaker_count": speaker_count, "text": speaker_text, "turns": speaker_turns, "text_raw": speaker_text_raw, "turns_raw": speaker_turns_raw, "text_syntax": speaker_text_syntax, "turns_syntax": speaker_turns_syntax, "syntax_postprocess_enabled": bool(syntax_postprocess_enabled), "syntax_postprocess_applied": bool(syntax_postprocess_applied), "syntax_sentence_count": int(syntax_sentence_count), "syntax_low_confidence_count": int(syntax_low_confidence_count), } if session_id: payload["session_id"] = session_id return payload def _diar_done_signal_payload( *, job_id: str, file_name: str, duration_s: float, model: str, speaker_count: int, speakers: list[str], syntax_postprocess_enabled: bool, syntax_postprocess_applied: bool, syntax_sentence_count: int, syntax_low_confidence_count: int, session_id: str | None, ) -> dict[str, Any]: payload: dict[str, Any] = { "type": "diar_done", "job_id": job_id, "file_name": file_name, "phase": "done", "duration_s": round(duration_s, 3), "model": model, "speaker_count": int(speaker_count), "speakers": list(speakers), "syntax_postprocess_enabled": bool(syntax_postprocess_enabled), "syntax_postprocess_applied": bool(syntax_postprocess_applied), "syntax_sentence_count": int(syntax_sentence_count), "syntax_low_confidence_count": int(syntax_low_confidence_count), } if session_id: payload["session_id"] = session_id return payload def _diar_text_done_signal_payload(payload: dict[str, Any]) -> dict[str, Any]: signal = { "type": "diar_text_done", "job_id": payload.get("job_id"), "file_name": payload.get("file_name"), "phase": payload.get("phase") or "done", "duration_s": payload.get("duration_s"), "speaker_count": payload.get("speaker_count"), "syntax_postprocess_enabled": payload.get("syntax_postprocess_enabled"), "syntax_postprocess_applied": payload.get("syntax_postprocess_applied"), "syntax_sentence_count": payload.get("syntax_sentence_count"), "syntax_low_confidence_count": payload.get("syntax_low_confidence_count"), } if payload.get("session_id"): signal["session_id"] = payload.get("session_id") return {key: value for key, value in signal.items() if value is not None} class FileJobWorkerService: def __init__(self, settings: FileModeSettings) -> None: self.service_name = _env("STACK_SERVICE_NAME", "stt_whisper_to_nats") self.nats_url = _env("NATS_URL") self._settings = settings self._diarization_runtime_label = _normalize_diarization_runtime_label(_env("DIARIZATION_RUNTIME_LABEL")) self._job_store = FileJobStore(settings.job_storage_dir) self._http: aiohttp.ClientSession | None = None self._nc: nats.NATS | None = None self._active_task: asyncio.Task | None = None self._running_job: str | None = None self._running_session_id: str | None = None self._lifecycle_lock = asyncio.Lock() self._cleanup_tasks: dict[str, asyncio.Task] = {} self._cancel_reasons: dict[str, str] = {} @property def active_job_id(self) -> str | None: return self._running_job @property def active_session_id(self) -> str | None: return self._running_session_id async def start(self) -> None: if self._http is None: timeout = aiohttp.ClientTimeout(total=None, connect=10.0) self._http = aiohttp.ClientSession(timeout=timeout) if self._nc is None and self.nats_url: self._nc = await nats.connect( servers=[self.nats_url], name=f"{self.service_name}-file-runner", max_reconnect_attempts=-1, reconnect_time_wait=2, ping_interval=10, ) logger.info("File runner connected to NATS %s", self.nats_url) elif self._nc is None: logger.warning("File runner started without NATS_URL; live file-mode signals are disabled") async def close(self) -> None: await self.cancel_active(reason="service shutdown") cleanup_tasks = list(self._cleanup_tasks.values()) self._cleanup_tasks.clear() for task in cleanup_tasks: task.cancel() if cleanup_tasks: await asyncio.gather(*cleanup_tasks, return_exceptions=True) if self._nc: try: await self._nc.drain() except Exception: await self._nc.close() self._nc = None if self._http: await self._http.close() self._http = None async def diarization_backend_status(self) -> tuple[bool, str]: await self.start() assert self._http is not None return await diarization_backend_ready(settings=self._settings, http=self._http) def _job_state_fields(self, payload: dict[str, Any]) -> dict[str, Any]: run_diarization = resolve_optional_bool( payload.get("run_diarization"), default=self._settings.diarization_enabled, ) return { "requested_diarization": resolve_optional_bool( payload.get("requested_diarization"), default=run_diarization, ), "run_diarization": run_diarization, "syntax_postprocess": bool(run_diarization) and resolve_optional_bool( payload.get("syntax_postprocess"), default=self._settings.diarization_syntax_postprocess_default, ), "diarization_backend_status": str(payload.get("diarization_backend_status") or "").strip(), "diarization_fallback_reason": str(payload.get("diarization_fallback_reason") or "").strip(), } async def start_job(self, payload: dict[str, Any]) -> None: job_id = str(payload.get("job_id") or "").strip() session_id = str(payload.get("session_id") or "").strip() or None file_name = str(payload.get("file_name") or "audio.bin").strip() or "audio.bin" if not job_id: raise ValueError("job_id is required") job_state_fields = self._job_state_fields(payload) await self.start() async with self._lifecycle_lock: await self._cancel_active_locked(reason="replaced by a new file job") lock_lease = await self._acquire_job_lock(payload=payload, job_state_fields=job_state_fields) try: await self._publish_job_event( job_id, { "type": "starting", "job_id": job_id, "file_name": file_name, "phase": "starting", **job_state_fields, }, session_id=session_id, ) task = asyncio.create_task( self._process_job_with_lock(payload, lock_lease=lock_lease), name=f"file-job-{job_id}", ) except Exception: if lock_lease is not None: await lock_lease.release() raise self._active_task = task self._running_job = job_id self._running_session_id = session_id task.add_done_callback(self._clear_active_task) async def _acquire_job_lock(self, *, payload: dict[str, Any], job_state_fields: dict[str, Any]) -> Any | None: if not job_lock_enabled(self._settings): return None assert self._http is not None job_id = str(payload.get("job_id") or "").strip() run_diarization = resolve_optional_bool(job_state_fields.get("run_diarization"), default=False) page_mode = str(payload.get("page_mode") or ("diarize" if run_diarization else "transcribe_only")).strip() step = "transcribe+diarize" if run_diarization else "transcribe" client = FileJobLockClient(settings=self._settings, http=self._http) lease = await client.acquire( owner=f"{self.service_name}:file:{job_id}", page_mode=page_mode, page_label_value=page_label(page_mode), job_id=job_id, session_id=str(payload.get("session_id") or "").strip(), file_name=str(payload.get("file_name") or "audio.bin").strip() or "audio.bin", step=step, ) if lease is not None: await lease.start_heartbeat() return lease async def cancel_active(self, *, session_id: str | None = None, reason: str = "cancel requested") -> str | None: async with self._lifecycle_lock: if session_id and self._running_session_id != session_id: return None return await self._cancel_active_locked(reason=reason) async def _cancel_active_locked(self, *, reason: str) -> str | None: job_id = str(self._running_job or "").strip() task = self._active_task if not job_id or task is None: return None self._job_store.request_cancel(job_id) self._cancel_reasons[job_id] = str(reason or "cancel requested").strip() or "cancel requested" logger.info("Cancelling active file job %s: %s", job_id, self._cancel_reasons[job_id]) if not task.done(): task.cancel() await asyncio.gather(task, return_exceptions=True) self._clear_active_task(task) self._job_store.delete_job(job_id) return job_id def _clear_active_task(self, task: asyncio.Task) -> None: if self._active_task is not task: return self._active_task = None self._running_job = None self._running_session_id = None async def _process_job_with_lock(self, payload: dict[str, Any], *, lock_lease: Any | None) -> None: try: await self._process_job(payload) finally: if lock_lease is not None: await lock_lease.release() async def _process_job(self, payload: dict[str, Any]) -> None: assert self._http is not None job_id = str(payload.get("job_id") or "").strip() file_name = str(payload.get("file_name") or "audio.bin").strip() or "audio.bin" file_path = str(payload.get("file_path") or "").strip() language = str(payload.get("language") or "").strip() or None session_id = str(payload.get("session_id") or "").strip() or None job_state_fields = self._job_state_fields(payload) run_diarization = bool(job_state_fields["run_diarization"]) syntax_postprocess = bool(job_state_fields["syntax_postprocess"]) if not job_id or not file_path: logger.warning("Dropping file job with missing job_id/file_path: %s", payload) return source_path = Path(file_path).resolve() if not source_path.exists(): await self._publish_job_event( job_id, { "type": "job_error", "job_id": job_id, "file_name": file_name, "phase": "prepare_media", "details": f"source file not found: {source_path}", }, session_id=session_id, ) if run_diarization: await self._publish_diar_error( job_id, file_name=file_name, details=f"source file not found: {source_path}", phase="prepare_media", session_id=session_id, ) self._schedule_job_cleanup(job_id) return diarization_task: asyncio.Task | None = None diarization_cancel_event: asyncio.Event | None = None diarization_requires_error = bool(run_diarization) diarization_config_error = "" job_done_published = False cleanup_job = False duration_s = 0.0 try: if self._job_store.is_cancel_requested(job_id): await self._publish_job_event( job_id, { "type": "job_error", "job_id": job_id, "file_name": file_name, "phase": "cancelled", "details": "cancel requested before processing started", }, session_id=session_id, ) if diarization_requires_error: await self._publish_diar_error( job_id, file_name=file_name, details="cancel requested before processing started", phase="cancelled", session_id=session_id, ) cleanup_job = True return if not self._settings.linto_http_url: await self._publish_job_event( job_id, { "type": "job_error", "job_id": job_id, "file_name": file_name, "phase": "prepare_media", "details": "LINTO_HTTP_URL is required for sequential file transcription", }, session_id=session_id, ) if diarization_requires_error: await self._publish_diar_error( job_id, file_name=file_name, details="LINTO_HTTP_URL is required for sequential file transcription", phase="prepare_media", session_id=session_id, ) cleanup_job = True return normalized_path = self._job_store.normalized_audio_path(job_id) await self._publish_job_event( job_id, { "type": "prepare_media", "job_id": job_id, "file_name": file_name, "phase": "prepare_audio", }, session_id=session_id, ) await asyncio.to_thread( write_normalized_audio_wav, source_path, output_path=normalized_path, sample_rate=16000, ) if self._job_store.is_cancel_requested(job_id): await self._publish_job_event( job_id, { "type": "job_error", "job_id": job_id, "file_name": file_name, "phase": "cancelled", "details": "cancel requested after audio preparation", }, session_id=session_id, ) if diarization_requires_error: await self._publish_diar_error( job_id, file_name=file_name, details="cancel requested after audio preparation", phase="cancelled", session_id=session_id, ) cleanup_job = True return duration_s = probe_wav_duration_s(normalized_path) chunk_total = estimate_chunk_total( duration_s=duration_s, chunk_seconds=self._settings.chunk_seconds, overlap_seconds=self._settings.chunk_overlap_seconds, ) await self._publish_job_event( job_id, { "type": "prepare_media", "job_id": job_id, "file_name": file_name, "phase": "prepare_media", "duration_s": round(duration_s, 2), "chunk_total": chunk_total, }, session_id=session_id, ) if run_diarization: if not self._settings.diarization_enabled: diarization_config_error = "file diarization is disabled" elif self._settings.diarization_http_url: diarization_artifact = await asyncio.to_thread( prepare_diarization_audio_artifact, source_path=source_path, normalized_path=normalized_path, diarization_audio_path=self._job_store.diarization_audio_path(job_id), requested_audio_format=self._settings.diarization_audio_format, wav_max_upload_mb=self._settings.diarization_wav_max_upload_mb, audio_max_upload_mb=self._settings.diarization_audio_max_upload_mb, ) diarization_input_path = diarization_artifact.path diarization_cancel_event = asyncio.Event() diarization_task = asyncio.create_task( self._run_diarization( job_id=job_id, upload_path=diarization_input_path, file_name=diarization_input_path.name, job_file_name=file_name, session_id=session_id, cancel_event=diarization_cancel_event, ), name=f"diarization-{job_id}", ) await self._publish_diar_progress( job_id, file_name=file_name, details="Speaker separation started. It runs in parallel with transcription.", phase="diarize_start", session_id=session_id, duration_s=duration_s, ) else: diarization_config_error = "DIARIZATION_HTTP_URL is required for file diarization" transcript = TimedTranscriptBuilder( soft_break_s=self._settings.text_soft_break_s, hard_break_s=self._settings.text_hard_break_s, require_word_timestamps=True, ) completed_chunks = 0 for chunk in iter_normalized_wav_chunk_files( normalized_path, output_dir=self._job_store.chunks_dir(job_id), sample_rate=16000, chunk_seconds=self._settings.chunk_seconds, overlap_seconds=self._settings.chunk_overlap_seconds, total_chunks=chunk_total, ): if self._job_store.is_cancel_requested(job_id): await self._publish_job_event( job_id, { "type": "job_error", "job_id": job_id, "file_name": file_name, "phase": "cancelled", "chunk_index": chunk.index, "chunk_total": chunk.total, "details": "cancel requested", }, session_id=session_id, ) if diarization_requires_error: await self._cancel_background_task(diarization_task, diarization_cancel_event) await self._publish_diar_error( job_id, file_name=file_name, details="cancel requested", phase="cancelled", session_id=session_id, duration_s=duration_s, ) cleanup_job = True return payload_json: dict[str, Any] | None = None last_error = "" for attempt in range(1, self._settings.chunk_retry_limit + 2): transcribe_task: asyncio.Task | None = None await self._publish_job_event( job_id, { "type": "chunk_started", "job_id": job_id, "file_name": file_name, "phase": "stream", "chunk_index": chunk.index, "chunk_total": chunk.total, "attempt": attempt, "chunk_start_s": round(chunk.start_s, 2), "chunk_end_s": round(chunk.end_s, 2), }, session_id=session_id, ) try: started_at = time.monotonic() transcribe_task = asyncio.create_task( transcribe_http_json( settings=self._settings, file_path=chunk.path, file_name=chunk.path.name, language=language, http=self._http, ) ) while True: try: payload_json = await asyncio.wait_for( asyncio.shield(transcribe_task), timeout=FILE_CHUNK_PROGRESS_INTERVAL_S, ) break except asyncio.TimeoutError: await self._publish_job_event( job_id, { "type": "chunk_progress", "job_id": job_id, "file_name": file_name, "phase": "transcribe_wait", "chunk_index": chunk.index, "chunk_total": chunk.total, "attempt": attempt, "chunk_start_s": round(chunk.start_s, 2), "chunk_end_s": round(chunk.end_s, 2), "elapsed_s": round(time.monotonic() - started_at, 1), }, session_id=session_id, ) break except Exception as exc: if transcribe_task is not None and not transcribe_task.done(): transcribe_task.cancel() try: await transcribe_task except BaseException: pass last_error = str(exc) if attempt > self._settings.chunk_retry_limit: break await self._publish_job_event( job_id, { "type": "chunk_retry", "job_id": job_id, "file_name": file_name, "phase": "retry", "chunk_index": chunk.index, "chunk_total": chunk.total, "attempt": attempt + 1, "details": str(exc), }, session_id=session_id, ) if payload_json is None: await self._publish_job_event( job_id, { "type": "job_error", "job_id": job_id, "file_name": file_name, "phase": "transcribe", "chunk_index": chunk.index, "chunk_total": chunk.total, "details": last_error or "chunk transcription failed", }, session_id=session_id, ) if diarization_requires_error: await self._cancel_background_task(diarization_task, diarization_cancel_event) await self._publish_diar_error( job_id, file_name=file_name, details=last_error or "chunk transcription failed", phase="transcribe", session_id=session_id, duration_s=duration_s, ) cleanup_job = True return try: transcript.append_chunk(payload=payload_json, chunk=chunk) except MissingWordTimestampsError as exc: await self._publish_job_event( job_id, { "type": "job_error", "job_id": job_id, "file_name": file_name, "phase": "merge", "chunk_index": chunk.index, "chunk_total": chunk.total, "chunk_start_s": round(chunk.start_s, 2), "chunk_end_s": round(chunk.end_s, 2), "details": str(exc), "debug": exc.debug, }, session_id=session_id, ) if diarization_requires_error: await self._cancel_background_task(diarization_task, diarization_cancel_event) await self._publish_diar_error( job_id, file_name=file_name, details=str(exc), phase="merge", session_id=session_id, duration_s=duration_s, ) cleanup_job = True return except Exception as exc: await self._publish_job_event( job_id, { "type": "job_error", "job_id": job_id, "file_name": file_name, "phase": "merge", "chunk_index": chunk.index, "chunk_total": chunk.total, "chunk_start_s": round(chunk.start_s, 2), "chunk_end_s": round(chunk.end_s, 2), "details": str(exc), }, session_id=session_id, ) if diarization_requires_error: await self._cancel_background_task(diarization_task, diarization_cancel_event) await self._publish_diar_error( job_id, file_name=file_name, details=str(exc), phase="merge", session_id=session_id, duration_s=duration_s, ) cleanup_job = True return completed_chunks += 1 current_text = transcript.render() self._job_store.write_transcript(job_id, current_text) await self._publish_job_event( job_id, { "type": "chunk_done", "job_id": job_id, "file_name": file_name, "phase": "stream", "chunk_index": chunk.index, "chunk_total": chunk.total, "chunks_completed": completed_chunks, "chunk_start_s": round(chunk.start_s, 2), "chunk_end_s": round(chunk.end_s, 2), "chars_total": len(current_text), }, session_id=session_id, ) try: chunk.path.unlink(missing_ok=True) except Exception: pass final_text = transcript.render() self._job_store.write_transcript(job_id, final_text) await self._publish_job_event( job_id, { "type": "job_done", "job_id": job_id, "file_name": file_name, "phase": "done", "chunk_total": max(chunk_total, completed_chunks or 1), "chunks_completed": completed_chunks, "chars_total": len(final_text), }, session_id=session_id, ) job_done_published = True if not run_diarization: cleanup_job = True return if diarization_config_error: await self._publish_diar_error( job_id, file_name=file_name, details=diarization_config_error, phase="diarize", session_id=session_id, duration_s=duration_s, ) cleanup_job = True return if diarization_task is None: await self._publish_diar_error( job_id, file_name=file_name, details="diarization task was not started", phase="diarize", session_id=session_id, duration_s=duration_s, ) cleanup_job = True return if not diarization_task.done(): await self._publish_diar_progress( job_id, file_name=file_name, details=_diarization_wait_details(self._diarization_runtime_label), phase="diarize_wait", session_id=session_id, duration_s=duration_s, ) diarization_payload = parse_diarization_result(await diarization_task) attributed_words = attribute_words( words=transcript.timed_words(), exclusive_segments=diarization_payload.exclusive_segments, ) raw_speaker_turns = build_speaker_turns( words=attributed_words, break_s=self._settings.diarization_turn_break_s, ) raw_speaker_text = render_speaker_text(turns=raw_speaker_turns) final_words = attributed_words final_speaker_turns = raw_speaker_turns final_speaker_text = raw_speaker_text syntax_speaker_turns: list[Any] = [] syntax_speaker_text = "" syntax_sentence_count = 0 syntax_low_confidence_count = 0 syntax_postprocess_applied = False if syntax_postprocess: syntax_result = apply_syntax_postprocess( words=attributed_words, sentence_gap_s=self._settings.diarization_syntax_sentence_gap_s, dominance_threshold=self._settings.diarization_syntax_dominance_threshold, turn_break_s=self._settings.diarization_turn_break_s, ) final_words = syntax_result.words final_speaker_turns = syntax_result.turns final_speaker_text = syntax_result.text syntax_speaker_turns = syntax_result.turns syntax_speaker_text = syntax_result.text syntax_sentence_count = syntax_result.sentence_count syntax_low_confidence_count = syntax_result.low_confidence_count syntax_postprocess_applied = True speaker_count = max(diarization_payload.speaker_count, len(diarization_payload.speakers)) diarization_result_payload: dict[str, Any] = { "type": "diar_done", "job_id": job_id, "file_name": file_name, "phase": "done", "duration_s": round(duration_s, 3), "model": diarization_payload.model, "speaker_count": speaker_count, "speakers": diarization_payload.speakers, "segments": diarization_segments_payload(diarization_payload.segments), "exclusive_segments": diarization_segments_payload(diarization_payload.exclusive_segments), "words": attributed_words_payload(final_words), "turns": speaker_turns_payload(final_speaker_turns), "words_raw": attributed_words_payload(attributed_words), "turns_raw": speaker_turns_payload(raw_speaker_turns), "text": final_speaker_text, "text_raw": raw_speaker_text, "text_syntax": syntax_speaker_text, "turns_syntax": speaker_turns_payload(syntax_speaker_turns), "syntax_postprocess_enabled": syntax_postprocess, "syntax_postprocess_applied": syntax_postprocess_applied, "syntax_sentence_count": syntax_sentence_count, "syntax_low_confidence_count": syntax_low_confidence_count, } if session_id: diarization_result_payload["session_id"] = session_id diar_done_signal = _diar_done_signal_payload( job_id=job_id, file_name=file_name, duration_s=duration_s, model=diarization_payload.model, speaker_count=speaker_count, speakers=diarization_payload.speakers, syntax_postprocess_enabled=syntax_postprocess, syntax_postprocess_applied=syntax_postprocess_applied, syntax_sentence_count=syntax_sentence_count, syntax_low_confidence_count=syntax_low_confidence_count, session_id=session_id, ) diar_text_payload = _diar_text_done_payload( job_id=job_id, file_name=file_name, duration_s=duration_s, speaker_count=speaker_count, speaker_text=final_speaker_text, speaker_turns=speaker_turns_payload(final_speaker_turns), speaker_text_raw=raw_speaker_text, speaker_turns_raw=speaker_turns_payload(raw_speaker_turns), speaker_text_syntax=syntax_speaker_text, speaker_turns_syntax=speaker_turns_payload(syntax_speaker_turns), syntax_postprocess_enabled=syntax_postprocess, syntax_postprocess_applied=syntax_postprocess_applied, syntax_sentence_count=syntax_sentence_count, syntax_low_confidence_count=syntax_low_confidence_count, session_id=session_id, ) await self._publish_diar_payload( job_id, diar_done_signal, store_payload=diarization_result_payload, ) await self._publish_diar_text_payload( job_id, _diar_text_done_signal_payload(diar_text_payload), store_payload=diar_text_payload, session_id=session_id, ) cleanup_job = True except asyncio.CancelledError: cancel_details = self._cancel_reasons.pop(job_id, None) or "cancel requested" logger.info( "File job cancelled (job_id=%s session_id=%s file_name=%s): %s", job_id, session_id or "-", file_name, cancel_details, ) await self._cancel_background_task(diarization_task, diarization_cancel_event) if not job_done_published: await self._publish_job_event( job_id, { "type": "job_error", "job_id": job_id, "file_name": file_name, "phase": "cancelled", "details": cancel_details, }, session_id=session_id, ) if diarization_requires_error: await self._publish_diar_error( job_id, file_name=file_name, details=cancel_details, phase="cancelled", session_id=session_id, duration_s=duration_s or None, ) cleanup_job = True return except Exception as exc: logger.exception("File job failed: %s", exc) if not job_done_published: await self._publish_job_event( job_id, { "type": "job_error", "job_id": job_id, "file_name": file_name, "phase": "worker_error", "details": str(exc), }, session_id=session_id, ) if diarization_requires_error: await self._publish_diar_error( job_id, file_name=file_name, details=str(exc), phase="worker_error", session_id=session_id, duration_s=duration_s or None, ) cleanup_job = True finally: await self._cancel_background_task(diarization_task, diarization_cancel_event) if cleanup_job: self._schedule_job_cleanup(job_id) async def _publish_job_event( self, job_id: str, payload: dict[str, Any], *, session_id: str | None = None, ) -> None: effective_session_id = str(session_id or payload.get("session_id") or "").strip() or None event_payload = dict(payload) if effective_session_id and not event_payload.get("session_id"): event_payload["session_id"] = effective_session_id event_payload = self._job_store.write_state(job_id, event_payload) if not self._nc: return encoded = json.dumps(event_payload, ensure_ascii=False).encode("utf-8") await self._nc.publish(self._settings.event_subject(job_id), encoded) if effective_session_id: await self._nc.publish(self._settings.session_event_subject(effective_session_id), encoded) async def _run_diarization( self, *, job_id: str, upload_path: Path, file_name: str, job_file_name: str, session_id: str | None, cancel_event: asyncio.Event, ) -> dict[str, Any]: assert self._http is not None last_error = "" for attempt in range(1, self._settings.diarization_retry_limit + 2): if cancel_event.is_set(): raise asyncio.CancelledError() try: logger.info( "Diarization attempt %s started (job_id=%s session_id=%s file_name=%s path=%s)", attempt, job_id, session_id or "-", job_file_name, upload_path, ) return await diarize_http_json( settings=self._settings, file_path=upload_path, file_name=file_name, job_id=job_id, session_id=session_id, job_file_name=job_file_name, http=self._http, cancel_event=cancel_event, ) except asyncio.CancelledError: logger.info( "Diarization task cancelled (job_id=%s session_id=%s file_name=%s)", job_id, session_id or "-", job_file_name, ) raise except Exception as exc: last_error = str(exc) logger.warning( "Diarization attempt %s failed (job_id=%s session_id=%s file_name=%s): %s", attempt, job_id, session_id or "-", job_file_name, exc, ) if attempt > self._settings.diarization_retry_limit: break raise RuntimeError(last_error or "diarization failed") @staticmethod async def _cancel_background_task(task: asyncio.Task | None, cancel_event: asyncio.Event | None = None) -> None: if cancel_event is not None: cancel_event.set() if task is None or task.done(): return logger.info("Cancelling background task %s", task.get_name()) task.cancel() try: await task except BaseException: pass def _schedule_job_cleanup(self, job_id: str, *, delay_s: float = FILE_JOB_TERMINAL_CLEANUP_DELAY_S) -> None: token = str(job_id or "").strip() if not token: return existing = self._cleanup_tasks.pop(token, None) if existing is not None: existing.cancel() task = asyncio.create_task(self._delete_job_after_delay(token, delay_s), name=f"cleanup-{token}") self._cleanup_tasks[token] = task def _clear(_task: asyncio.Task) -> None: if self._cleanup_tasks.get(token) is _task: self._cleanup_tasks.pop(token, None) task.add_done_callback(_clear) async def _delete_job_after_delay(self, job_id: str, delay_s: float) -> None: if delay_s > 0: await asyncio.sleep(delay_s) self._job_store.delete_job(job_id) async def _publish_diar_payload( self, job_id: str, payload: dict[str, Any], *, store_payload: dict[str, Any] | None = None, ) -> None: self._job_store.write_diarization_state(job_id, payload) if store_payload is not None: self._job_store.write_diarization_payload(job_id, store_payload) if not self._nc: return encoded = json.dumps(payload, ensure_ascii=False).encode("utf-8") await self._nc.publish(self._settings.diar_event_subject(job_id), encoded) async def _publish_diar_text_payload( self, job_id: str, payload: dict[str, Any], *, session_id: str | None = None, store_payload: dict[str, Any] | None = None, ) -> None: self._job_store.write_diarization_text(job_id, store_payload or payload) if not self._nc: return encoded = json.dumps(payload, ensure_ascii=False).encode("utf-8") await self._nc.publish(self._settings.diar_text_event_subject(job_id), encoded) effective_session_id = str(session_id or payload.get("session_id") or "").strip() or None if effective_session_id: await self._nc.publish(self._settings.diar_text_session_subject(effective_session_id), encoded) async def _publish_diar_progress( self, job_id: str, *, file_name: str, details: str, phase: str, session_id: str | None = None, duration_s: float | None = None, ) -> None: base_payload: dict[str, Any] = { "type": "diar_progress", "job_id": job_id, "file_name": file_name, "phase": phase, "details": str(details or "").strip() or "speaker separation is in progress", } if duration_s is not None: base_payload["duration_s"] = round(float(duration_s), 3) if session_id: base_payload["session_id"] = session_id await self._publish_diar_payload(job_id, base_payload) await self._publish_diar_text_payload( job_id, { **base_payload, "type": "diar_text_progress", }, session_id=session_id, ) async def _publish_diar_error( self, job_id: str, *, file_name: str, details: str, phase: str, session_id: str | None = None, duration_s: float | None = None, ) -> None: base_payload: dict[str, Any] = { "type": "diar_error", "job_id": job_id, "file_name": file_name, "phase": phase, "details": str(details or "").strip() or "unknown diarization error", } if duration_s is not None: base_payload["duration_s"] = round(float(duration_s), 3) if session_id: base_payload["session_id"] = session_id await self._publish_diar_payload(job_id, base_payload) await self._publish_diar_text_payload( job_id, { **base_payload, "type": "diar_text_error", }, session_id=session_id, ) def configure_logging() -> None: level = _env("LOG_LEVEL", "INFO").upper() logging.basicConfig( level=getattr(logging, level, logging.INFO), format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", ) async def _main_async() -> None: stream_bridge_enabled = _env_bool("STREAM_BRIDGE_ENABLED", True) webui_enabled = _env_bool("WEBUI_ENABLED", True) file_mode_settings = load_file_mode_settings() if not any((stream_bridge_enabled, webui_enabled)): raise ValueError("At least one runtime must be enabled: WEBUI_ENABLED or STREAM_BRIDGE_ENABLED") bridge_service = SttWhisperToNatsService(file_mode_settings) if stream_bridge_enabled else None tasks: list[asyncio.Task] = [] stop_event = asyncio.Event() def request_stop() -> None: if stop_event.is_set(): return stop_event.set() if bridge_service is not None: bridge_service.stop() loop = asyncio.get_running_loop() for sig in (signal.SIGINT, signal.SIGTERM): try: loop.add_signal_handler(sig, request_stop) except NotImplementedError: pass if webui_enabled: from src.webui.server import run_webui host = _env("WEBUI_HOST", "0.0.0.0") port = int(_env("WEBUI_PORT", "8090") or "8090") base_path = _env("WEBUI_BASE_PATH", "") tasks.append(asyncio.create_task(run_webui(host=host, port=port, base_path=base_path), name="webui")) if bridge_service is not None: tasks.append(asyncio.create_task(bridge_service.run(), name="stream-bridge")) stopper = asyncio.create_task(stop_event.wait(), name="stop-wait") try: done, pending = await asyncio.wait(tasks + [stopper], return_when=asyncio.FIRST_COMPLETED) if stopper not in done: for task in done: if task.cancelled(): continue exc = task.exception() if exc is not None: logger.exception("Runtime task failed: %s", exc) else: logger.warning("Runtime task exited unexpectedly: %s", task.get_name()) request_stop() await asyncio.wait(tasks, timeout=5.0) finally: request_stop() stopper.cancel() for task in tasks: if task.get_name() == "webui": task.cancel() await asyncio.gather(*tasks, return_exceptions=True) await asyncio.gather(stopper, return_exceptions=True) def main() -> None: configure_logging() try: asyncio.run(_main_async()) except KeyboardInterrupt: pass except Exception as exc: logger.exception("Service terminated with error: %s", exc) raise SystemExit(1) from exc if __name__ == "__main__": main()