voice-agent/lxc-frontend/app/services/pipeline.py

import asyncio
import json
import logging
from datetime import datetime, timezone
from pathlib import Path
from typing import Awaitable, TypeVar

from sqlmodel import Session

from app.core.config import settings
from app.core.db import engine
from app.core.mac_client import MacClient
from app.models.job import Job, JobStatus

log = logging.getLogger("voice-agent.pipeline")

T = TypeVar("T")

HEARTBEAT_INTERVAL = 3.0  # Sekunden


def _update(job_id: int, **fields) -> None:
    with Session(engine) as s:
        job = s.get(Job, job_id)
        if not job:
            return
        for k, v in fields.items():
            setattr(job, k, v)
        job.updated_at = datetime.now(timezone.utc)
        s.add(job)
        s.commit()


def _existing(path_str: str) -> Path | None:
    if not path_str:
        return None
    p = Path(path_str)
    return p if p.exists() and p.stat().st_size > 0 else None


def _start_step(job_id: int, status: str, progress: int) -> None:
    """Setzt status/progress + step_started_at + initialen Heartbeat."""
    now = datetime.now(timezone.utc)
    _update(
        job_id,
        status=status,
        progress=progress,
        step_started_at=now,
        last_heartbeat_at=now,
        error="",
    )


async def _with_heartbeat(job_id: int, coro: Awaitable[T]) -> T:
    """Führt `coro` aus und pingt parallel alle paar Sekunden last_heartbeat_at.

    Damit sieht die UI live, ob ein Schritt noch lebt — auch wenn der Mac
    minutenlang an einer Ollama-Generierung arbeitet.
    """
    stop = asyncio.Event()

    async def beat() -> None:
        while not stop.is_set():
            try:
                _update(job_id, last_heartbeat_at=datetime.now(timezone.utc))
            except Exception:  # noqa: BLE001
                pass  # Heartbeat darf den Job nie killen
            try:
                await asyncio.wait_for(stop.wait(), timeout=HEARTBEAT_INTERVAL)
            except asyncio.TimeoutError:
                pass

    beat_task = asyncio.create_task(beat())
    try:
        return await coro
    finally:
        stop.set()
        try:
            await beat_task
        except Exception:  # noqa: BLE001
            pass


async def run_pipeline(job_id: int) -> None:
    """Pipeline: transcribe → preload → summarize → protocol → DOCX → PDF.

    Idempotent: bereits vorhandene Zwischenergebnisse werden weiterverwendet,
    damit ein 'Erneut versuchen' nicht die teure Transkription wiederholt.
    """
    client = MacClient()

    with Session(engine) as s:
        job = s.get(Job, job_id)
        if not job:
            log.error("Job %s not found", job_id)
            return
        audio_path = settings.upload_dir / job.filename
        title = job.title or audio_path.stem
        profile = job.profile or "meeting"
        existing_tx_path  = _existing(job.transcript_path)
        existing_sum_path = _existing(job.summary_path)
        existing_pro_path = _existing(job.protocol_path)
        existing_dx_path  = _existing(job.docx_path)
        existing_pdf_path = _existing(job.pdf_path)

    job_dir = settings.result_dir / str(job_id)
    job_dir.mkdir(parents=True, exist_ok=True)

    try:
        # 1) TRANSKRIPTION
        if existing_tx_path:
            log.info("[job %s] reuse transcript %s", job_id, existing_tx_path.name)
            tx = json.loads(existing_tx_path.read_text(encoding="utf-8"))
            _update(job_id, status=JobStatus.TRANSCRIBING, progress=45)
        else:
            _start_step(job_id, JobStatus.TRANSCRIBING, 10)
            log.info("[job %s] transcribe (profile=%s)", job_id, profile)
            tx = await _with_heartbeat(job_id, client.transcribe(audio_path, language="de", job_id=job_id))
            transcript_path = job_dir / "transcript.json"
            transcript_path.write_text(json.dumps(tx, ensure_ascii=False, indent=2), encoding="utf-8")
            _update(job_id, transcript_path=str(transcript_path), progress=45)

        transcript_text = tx.get("text") or "\n".join(
            s.get("text", "") for s in tx.get("segments", [])
        )

        # 2) SUMMARY — vorher Preload, damit Ollama nicht mitten drin lädt
        if existing_sum_path:
            log.info("[job %s] reuse summary %s", job_id, existing_sum_path.name)
            summary = json.loads(existing_sum_path.read_text(encoding="utf-8"))
            _update(job_id, status=JobStatus.SUMMARIZING, progress=70)
        else:
            _start_step(job_id, JobStatus.SUMMARIZING, 50)
            log.info("[job %s] preload model", job_id)
            try:
                await _with_heartbeat(job_id, client.preload(profile=profile, job_id=job_id))
            except Exception as e:  # noqa: BLE001
                log.warning("[job %s] preload failed (non-fatal): %s", job_id, e)
            _update(job_id, progress=55)
            log.info("[job %s] summarize", job_id)
            summary = await _with_heartbeat(
                job_id,
                client.summarize(transcript_text, title=title, profile=profile, job_id=job_id),
            )
            summary_path = job_dir / "summary.json"
            summary_path.write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
            _update(job_id, summary_path=str(summary_path), progress=70)

        # 3) PROTOKOLL
        if existing_pro_path:
            log.info("[job %s] reuse protocol %s", job_id, existing_pro_path.name)
            protocol = json.loads(existing_pro_path.read_text(encoding="utf-8"))
            _update(job_id, status=JobStatus.PROTOCOLLING, progress=85)
        else:
            _start_step(job_id, JobStatus.PROTOCOLLING, 75)
            log.info("[job %s] protocol", job_id)
            protocol = await _with_heartbeat(
                job_id,
                client.protocol(transcript_text, summary, title=title, profile=profile, job_id=job_id),
            )
            protocol_path = job_dir / "protocol.json"
            protocol_path.write_text(json.dumps(protocol, ensure_ascii=False, indent=2), encoding="utf-8")
            _update(job_id, protocol_path=str(protocol_path), progress=85)

        # 4) DOCX
        if existing_dx_path:
            log.info("[job %s] reuse docx %s", job_id, existing_dx_path.name)
        else:
            _start_step(job_id, JobStatus.EXPORTING, 90)
            log.info("[job %s] export docx", job_id)
            docx_bytes = await _with_heartbeat(
                job_id,
                client.export_docx(protocol, profile=profile, job_id=job_id),
            )
            docx_path = job_dir / "protocol.docx"
            docx_path.write_bytes(docx_bytes)
            _update(job_id, docx_path=str(docx_path), progress=95)

        # 5) PDF
        if existing_pdf_path:
            log.info("[job %s] reuse pdf %s", job_id, existing_pdf_path.name)
            _update(job_id, status=JobStatus.DONE, progress=100, error="")
        else:
            _update(job_id, status=JobStatus.EXPORTING, progress=96)
            log.info("[job %s] export pdf", job_id)
            pdf_bytes = await _with_heartbeat(
                job_id,
                client.export_pdf(protocol, profile=profile, job_id=job_id),
            )
            pdf_path = job_dir / "protocol.pdf"
            pdf_path.write_bytes(pdf_bytes)
            _update(job_id, pdf_path=str(pdf_path), status=JobStatus.DONE, progress=100, error="")
        log.info("[job %s] done", job_id)
    except Exception as e:  # noqa: BLE001
        log.exception("[job %s] failed", job_id)
        _update(job_id, status=JobStatus.FAILED, error=f"{type(e).__name__}: {e}")