tueit_Transkriptor/api/pipeline.py

import asyncio
import logging
import os
import tempfile
import traceback
from datetime import datetime

import pipewire
from api.state import state, Status
from api.router import broadcast
from config import load as load_config
from transcription import engine as transcription_engine
from llm import OllamaClient
from output import write_solo_docs, write_meeting_docs

logger = logging.getLogger(__name__)


async def run_pipeline():
    cfg = load_config()
    recorder = getattr(state, "_recorder", None)
    if recorder is None:
        return

    output_dir = getattr(state, "_recording_output_dir", cfg["output"]["path"])
    instructions = getattr(state, "_recording_instructions", "")
    diar_cfg = cfg.get("diarization", {})
    use_diarization = diar_cfg.get("enabled") and diar_cfg.get("hf_token")

    # Drop the capture chain as soon as the audio is captured — transcription
    # runs for minutes and must not keep the PipeWire clock domains coupled.
    try:
        recorder.stop()
    finally:
        pipewire.teardown()

    await state.set_status(Status.PROCESSING)
    await broadcast({"event": "processing"})

    wav_path = None
    try:
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
            wav_path = f.name
        recorder.save_wav(wav_path)

        if use_diarization:
            await _run_meeting_pipeline(cfg, wav_path, output_dir, instructions, diar_cfg)
        else:
            await _run_solo_pipeline(cfg, wav_path, output_dir, instructions)

    except Exception as e:
        tb = traceback.format_exc()
        logger.error("Pipeline error:\n%s", tb)
        state.last_error = str(e)
        await state.set_status(Status.ERROR)
        await broadcast({"event": "error", "message": str(e)})
    finally:
        state.recording_user = None
        state._recording_output_dir = None
        state._recording_instructions = ""
        state._speakers_event = None
        state._pending_aligned_segments = None
        state._speaker_names = None
        if wav_path:
            try:
                os.unlink(wav_path)
            except OSError:
                pass


async def _run_solo_pipeline(cfg, wav_path, output_dir, instructions):
    """Original single-document pipeline (no diarization)."""
    raw_text = await transcription_engine.transcribe_file(
        wav_path,
        language=cfg["whisper"]["language"],
        model_name=cfg["whisper"]["model"],
        device=cfg["whisper"]["device"],
        base_url=cfg["whisper"].get("base_url", ""),
        backend=cfg["whisper"].get("backend", "openai"),
    )
    await broadcast({"event": "transcribed", "raw": raw_text})

    client = OllamaClient(base_url=cfg["ollama"]["base_url"])
    punctuated, refined = await asyncio.gather(
        client.punctuate(raw_text, model=cfg["ollama"]["model"]),
        client.refine(raw_text=raw_text, instructions=instructions, model=cfg["ollama"]["model"]),
    )
    title, tldr = await client.generate_title_and_tldr(refined, model=cfg["ollama"]["model"])

    dt = datetime.now()
    paths = write_solo_docs(raw_text=punctuated, refined=refined, output_dir=output_dir, dt=dt, title=title, tldr=tldr)

    await state.set_status(Status.IDLE)
    await broadcast({
        "event": "saved",
        "path": paths["index"],
        "title": title,
        "paths": paths,
    })


async def _run_meeting_pipeline(cfg, wav_path, output_dir, instructions, diar_cfg):
    """Diarization pipeline: 3 documents, speaker identification."""
    from diarization import Diarizer
    from alignment import align_segments

    diarizer = Diarizer(hf_token=diar_cfg["hf_token"])
    whisper_task = asyncio.create_task(
        transcription_engine.transcribe_file(
            wav_path,
            language=cfg["whisper"]["language"],
            model_name=cfg["whisper"]["model"],
            device=cfg["whisper"]["device"],
            base_url=cfg["whisper"].get("base_url", ""),
            backend=cfg["whisper"].get("backend", "openai"),
            with_segments=True,
        )
    )
    diar_task = asyncio.create_task(diarizer.diarize(wav_path))
    whisper_segs, speaker_segs = await asyncio.gather(whisper_task, diar_task)

    aligned = align_segments(whisper_segs, speaker_segs)
    await broadcast({"event": "transcribed", "raw": " ".join(t for _, t in aligned)})

    excerpt = "\n".join(f"{s}: {t}" for s, t in aligned[:20])
    client = OllamaClient(base_url=cfg["ollama"]["base_url"])
    name_map = await client.identify_speakers(excerpt, model=cfg["ollama"]["model"])

    if not name_map:
        excerpts_per_speaker = _build_excerpts(aligned)
        state._speakers_event = asyncio.Event()
        state._pending_aligned_segments = aligned
        await state.set_status(Status.AWAITING_SPEAKERS)
        await broadcast({"event": "speakers_unknown", "speakers": [
            {"id": spk, "excerpts": exs}
            for spk, exs in excerpts_per_speaker.items()
        ]})
        await state._speakers_event.wait()
        name_map = state._speaker_names or {}

    def resolve(label):
        name = name_map.get(label, "")
        if name:
            return name
        num = label.replace("SPEAKER_", "").lstrip("0") or "1"
        return f"Sprecher {num}"

    named_aligned = [(resolve(spk), text) for spk, text in aligned]
    speakers = sorted({spk for spk, _ in named_aligned})

    total_secs = sum(s["end"] - s["start"] for s in whisper_segs) if whisper_segs else 0
    duration_min = max(1, round(total_secs / 60))

    raw_transcript = "\n\n".join(f"**{spk}:** {txt}" for spk, txt in named_aligned)
    summary, punctuated_transcript = await asyncio.gather(
        client.summarize(raw_transcript, model=cfg["ollama"]["model"]),
        client.punctuate(raw_transcript, model=cfg["ollama"]["model"]),
    )
    title, tldr = await client.generate_title_and_tldr(summary, model=cfg["ollama"]["model"])

    dt = datetime.now()
    paths = write_meeting_docs(
        aligned_segments=named_aligned,
        summary=summary,
        speakers=speakers,
        duration_min=duration_min,
        output_dir=output_dir,
        dt=dt,
        title=title,
        tldr=tldr,
        transcript_text=punctuated_transcript,
    )

    await state.set_status(Status.IDLE)
    await broadcast({
        "event": "saved",
        "path": paths["index"],
        "title": title,
        "meeting": True,
        "paths": paths,
    })


def _build_excerpts(aligned: list[tuple[str, str]], max_per_speaker: int = 4) -> dict[str, list[str]]:
    """Build a dict of speaker → list of text excerpts."""
    from collections import defaultdict
    buckets: dict[str, list[str]] = defaultdict(list)
    for spk, text in aligned:
        if len(buckets[spk]) < max_per_speaker:
            buckets[spk].append(text[:200])
    return dict(buckets)