diarization_component.py

from __future__ import annotations

from typing import Any

try:
    from .config import VoiceRuntimeConfig
    from .diarization import (
        attach_speakers_to_segments,
        attach_speakers_to_words,
        build_diarization_summary,
        run_diarization,
    )
except ImportError:  # HF flat-root execution fallback
    from config import VoiceRuntimeConfig
    from diarization import (
        attach_speakers_to_segments,
        attach_speakers_to_words,
        build_diarization_summary,
        run_diarization,
    )


def apply_diarization(
    wav_path: str,
    config: VoiceRuntimeConfig,
    sample_rate: int,
    segments: list[dict[str, Any]],
    words: list[dict[str, Any]],
) -> tuple[list[dict[str, Any]], dict[str, Any]]:
    """Run diarization and project speaker labels onto words and segments."""
    diarization_segments = run_diarization(
        wav_path=wav_path,
        config=config,
        sample_rate=sample_rate,
    )

    attach_speakers_to_words(words=words, diarization_segments=diarization_segments)
    attach_speakers_to_segments(segments=segments, words=words, sample_rate=sample_rate)
    summary = build_diarization_summary(diarization_segments)
    return diarization_segments, summary


def run_diarization_only(
    wav_path: str,
    config: VoiceRuntimeConfig,
    sample_rate: int,
) -> tuple[list[dict[str, Any]], dict[str, Any]]:
    """Run diarization without any transcription/alignment dependencies."""
    diarization_segments = run_diarization(
        wav_path=wav_path,
        config=config,
        sample_rate=sample_rate,
    )
    summary = build_diarization_summary(diarization_segments)
    return diarization_segments, summary