my_voice_assistant_v3_jamulix/app/api/ws.py

"""WebSocket-Echtzeit-Chat und -Sprache.

- /ws/chat : Text rein (JSON pro Turn), Antwort als Event-Folge zurueck.
- /ws/voice: Audio rein (binaere Chunks + Control), Transkription -> selbe Pipeline.

Antwort-Events: ack -> [token*] -> [audio*] -> semantic -> done.
Mit {"stream":true} kommen LLM-Token live, mit {"audio_stream":true} das Audio
satzweise (chunked TTS). /ws/voice sendet zuvor ein transcript-Event.

Barge-in: Ein {"type":"interrupt"}-Frame oder eine neue Eingabe bricht eine laufende
Antwort ab (-> interrupted-Event). Der Antwort-Turn laeuft als abbrechbarer Task.

Spaeter (eigene Increments): echte partielle Live-Transkripte (Streaming-STT-Dienst),
WebRTC.
"""

import asyncio
import json

from fastapi import APIRouter, WebSocket, WebSocketDisconnect

from app.config import settings
from app.errors import RoutingError
from app.dependencies import (
    get_store,
    resolve_route,
    build_orchestrator,
    resolve_output_endpoint,
    piper_voice_for_language,
    voice_for_route,
)
from app.store import ANONYMOUS_USER_ID, SessionOwnershipError
from app.audio.vad import EnergyVAD
from app.core.memory_extractor import maybe_schedule_extraction
from app.quota import (enforce_quota, record_usage, QuotaExceededError,
                       web_search_within_budget, record_web_search)
from app.auth import authenticate, _cookie_value, CAPABILITY_COOKIE

router = APIRouter()

_OVERRIDE_KEYS = (
    "input_endpoint",
    "output_endpoint",
    "language",
    "stt_provider",
    "llm_provider",
    "tts_provider",
)


def _capability_token(websocket: WebSocket, token: str | None) -> str | None:
    """Token aus ?token=, ?k= oder dem va_token-Cookie (Senioren-Login)."""
    return (
        token
        or websocket.query_params.get("k")
        or _cookie_value(websocket.headers.get("cookie"), CAPABILITY_COOKIE)
    )


def _authenticate(websocket: WebSocket, token: str | None):
    # Forward-Auth (SSO) greift auch beim WS-Handshake: SSOwat injiziert den
    # Identitaets-Header in den Upgrade-Request -> aus websocket.headers lesbar.
    client_host = websocket.client.host if websocket.client else ""
    return authenticate(websocket.headers, client_host, _capability_token(websocket, token))


async def _resolve(user, session_id, options, store=None):
    """Loest Route + Orchestrator + Output-Endpunkt auf (kann RoutingError/Ownership werfen)."""
    overrides = {key: options.get(key) for key in _OVERRIDE_KEYS}
    # Web-Such-Tagesbudget aufgebraucht -> Web-Suche fuer diesen Turn aus.
    if store is not None and not web_search_within_budget(user, store):
        overrides["web_search_enabled"] = False
    route = resolve_route(user, session_id, overrides)
    orchestrator = build_orchestrator(route)
    output = await resolve_output_endpoint(route)
    return route, orchestrator, output


async def _run_turn(
    websocket, store, user, session_id, route, orchestrator, output, text, options,
    effective_voice: str | None = None,
):
    """Faehrt einen Antwort-Turn und streamt die Events an den Client."""
    conversation = (
        store.get_recent_messages(session_id, settings.history_max_messages)
        if session_id
        else []
    )
    memories = store.get_memories(user.id)
    llm_context = list(conversation)
    user_context_parts = []
    if user.id != ANONYMOUS_USER_ID:
        user_context_parts.append(f"Du sprichst mit {user.display_name}.")
    if memories:
        user_context_parts.append(
            "Was du ueber den Nutzer weisst:\n" + "\n".join(f"- {m.content}" for m in memories)
        )
    if user_context_parts:
        llm_context = [{"role": "system", "content": "\n".join(user_context_parts)}] + llm_context

    try:
        enforce_quota(user, store)
    except QuotaExceededError as exc:
        await websocket.send_json({"type": "error", "status": 429, "detail": str(exc)})
        return

    await websocket.send_json({"type": "ack", "route": route.as_dict()})

    # Stimme: explizit angeforderte gewinnt, sonst die vom Caller (Sprach-Turn)
    # vorberechnete sprachpassende Stimme, sonst folgt sie der Routensprache
    # (greift v. a. beim Text-Chat ohne Spracherkennung).
    explicit_voice = options.get("voice")
    if explicit_voice:
        voice = explicit_voice
    elif effective_voice is not None:
        voice = effective_voice
    else:
        voice = voice_for_route(route.tts_provider, route.language, route.voice_gender)
    stream = bool(options.get("stream"))
    # text_only: Geräte-TTS (Web Speech API) spricht selbst -> kein Server-Audio erzeugen/senden.
    text_only = bool(options.get("text_only"))
    # audio_stream: explizite Anfrage gewinnt, sonst der serverseitige Default (Admin).
    audio_stream = False if text_only else (
        bool(options["audio_stream"]) if "audio_stream" in options
        else settings.audio_stream_default
    )

    on_token = None
    if stream:
        async def on_token(delta):
            await websocket.send_json({"type": "token", "text": delta})

    on_audio = None
    if audio_stream:
        audio_seq = 0

        async def on_audio(chunk):
            nonlocal audio_seq
            await websocket.send_json({"type": "audio", "seq": audio_seq})
            audio_seq += 1
            await websocket.send_bytes(chunk)

    # Filler (Beruhigungs-/Geduldssätze) als eigener Event: Status-Zeile + Geräte-TTS.
    on_filler = None
    if stream or audio_stream:
        async def on_filler(filler_text):
            await websocket.send_json({"type": "filler", "text": filler_text})

    try:
        if stream or audio_stream:
            trace, audio = await orchestrator.chat_stream(
                text,
                language=route.language,
                voice=voice,
                output=output,
                history=llm_context,
                on_token=on_token,
                on_audio=on_audio,
                on_filler=on_filler,
                text_only=text_only,
            )
        else:
            trace, audio = await orchestrator.chat_text(
                text,
                language=route.language,
                voice=voice,
                output=output,
                history=llm_context,
                text_only=text_only,
            )
    except Exception as exc:
        await websocket.send_json({"type": "error", "status": 502, "detail": str(exc)})
        return

    record_usage(user, store, len(text) + len(trace.semantic_response or ""))
    record_web_search(user, store, trace.web_searches)

    if session_id:
        store.append_message(session_id, user.id, "user", text)
        store.append_message(session_id, user.id, "assistant", trace.semantic_response)
        maybe_schedule_extraction(store, user.id, session_id)

    await websocket.send_json(
        {"type": "semantic", "text": trace.semantic_response,
         "spoken": trace.spoken_response, "citations": trace.citations}
    )
    # Im text_only-Modus kommt kein Audio (Gerät spricht selbst).
    if not audio_stream and not text_only:
        await websocket.send_bytes(audio)
    await websocket.send_json({"type": "done", "audio_format": "pcm", "sample_rate": 24000})


async def _cancel_active(task, websocket) -> None:
    """Bricht einen laufenden Antwort-Turn ab (Barge-in) und meldet 'interrupted'."""
    if task is None or task.done():
        return
    task.cancel()
    try:
        await task
    except asyncio.CancelledError:
        pass
    await websocket.send_json({"type": "interrupted"})


async def _chat_turn(websocket, store, user, session_id, text, options):
    try:
        route, orchestrator, output = await _resolve(user, session_id, options, store)
    except SessionOwnershipError as exc:
        await websocket.send_json({"type": "error", "status": 403, "detail": str(exc)})
        return
    except RoutingError as exc:
        await websocket.send_json({"type": "error", "status": 422, "detail": str(exc)})
        return
    await _run_turn(websocket, store, user, session_id, route, orchestrator, output, text, options)


async def _voice_turn(websocket, store, user, session_id, audio, fmt, options):
    try:
        route, orchestrator, output = await _resolve(user, session_id, options, store)
    except SessionOwnershipError as exc:
        await websocket.send_json({"type": "error", "status": 403, "detail": str(exc)})
        return
    except RoutingError as exc:
        await websocket.send_json({"type": "error", "status": 422, "detail": str(exc)})
        return

    try:
        if settings.stt_force_language:
            # Feste Zielsprache statt Auto-Detect -> robuster bei kurzen/verrauschten
            # (Handy-)Clips; verhindert "falsche Sprache erkannt = Blödsinn".
            transcript, _ = await orchestrator.stt.transcribe_detect(
                audio, fmt=fmt, language=route.language
            )
        else:
            # Gesprochene Sprache automatisch erkennen; weicht sie von der Zielsprache
            # ab, wird in die Zielsprache übersetzt (Fremdsprachen-Diktat).
            transcript, detected = await orchestrator.stt.transcribe_detect(
                audio, fmt=fmt, language=None
            )
            if (transcript and transcript.strip() and detected
                    and detected.lower() != (route.language or "").lower()):
                transcript = await orchestrator.translate(transcript, route.language)
        effective_voice = voice_for_route(route.tts_provider, route.language, route.voice_gender)
    except Exception as exc:
        await websocket.send_json({"type": "error", "status": 502, "detail": str(exc)})
        return

    await websocket.send_json({
        "type": "transcript",
        "text": transcript,
    })
    if not transcript or not transcript.strip():
        await websocket.send_json({
            "type": "error",
            "detail": "Keine Sprache erkannt — bitte erneut sprechen.",
        })
        return
    await _run_turn(
        websocket, store, user, session_id, route, orchestrator, output, transcript, options,
        effective_voice=effective_voice,
    )


@router.websocket("/ws/chat")
async def ws_chat(websocket: WebSocket, session_id: str | None = None, token: str | None = None):
    user = _authenticate(websocket, token)
    if user is None:
        await websocket.close(code=1008)
        return
    await websocket.accept()
    store = get_store()
    active = None

    try:
        while True:
            msg = await websocket.receive_json()
            if msg.get("type") == "interrupt":
                await _cancel_active(active, websocket)
                active = None
                continue
            text = (msg.get("text") or "").strip()
            if not text:
                await websocket.send_json({"type": "error", "detail": "empty text"})
                continue
            await _cancel_active(active, websocket)  # Barge-in bei neuer Eingabe
            active = asyncio.create_task(
                _chat_turn(websocket, store, user, session_id, text, msg)
            )
    except WebSocketDisconnect:
        if active and not active.done():
            active.cancel()
        return


@router.websocket("/ws/voice")
async def ws_voice(websocket: WebSocket, session_id: str | None = None, token: str | None = None):
    user = _authenticate(websocket, token)
    if user is None:
        await websocket.close(code=1008)
        return
    await websocket.accept()
    store = get_store()

    audio_buffer = bytearray()
    fmt = "wav"
    active = None
    vad = None
    start_options: dict = {}  # Konfig aus dem start-Frame (Provider, audio_stream, language)

    async def _start_voice(audio: bytes, options: dict):
        nonlocal active
        await _cancel_active(active, websocket)  # Barge-in bei neuer Aeusserung
        active = asyncio.create_task(
            _voice_turn(websocket, store, user, session_id, audio, fmt, options)
        )

    try:
        while True:
            message = await websocket.receive()
            if message["type"] == "websocket.disconnect":
                if active and not active.done():
                    active.cancel()
                return

            if message.get("bytes") is not None:
                audio_buffer.extend(message["bytes"])
                # VAD: Aeusserungsende automatisch erkennen (opt-in via start-Frame).
                if vad is not None and vad.feed(message["bytes"]):
                    audio = bytes(audio_buffer)
                    audio_buffer.clear()
                    vad.reset()
                    await _start_voice(audio, start_options)
                continue

            raw = message.get("text")
            if raw is None:
                continue
            try:
                control = json.loads(raw)
            except ValueError:
                await websocket.send_json({"type": "error", "detail": "invalid control frame"})
                continue

            ctype = control.get("type")
            if ctype == "interrupt":
                await _cancel_active(active, websocket)
                active = None
                continue
            if ctype == "start":
                audio_buffer.clear()
                start_options = control  # Provider/audio_stream/language fuer den Turn merken
                fmt = control.get("format", "wav")
                if control.get("vad"):
                    vad = EnergyVAD(
                        sample_rate=control.get("sample_rate", 16000),
                        threshold=control.get("vad_threshold", 500.0),
                        silence_ms=control.get("vad_silence_ms", 700.0),
                    )
                else:
                    vad = None
                continue
            if ctype != "end":
                continue

            if not audio_buffer:
                await websocket.send_json({"type": "error", "detail": "no audio received"})
                continue

            audio = bytes(audio_buffer)
            audio_buffer.clear()
            if vad is not None:
                vad.reset()
            # start-Frame-Konfig + end-Frame zusammenfuehren (end kann ueberschreiben)
            await _start_voice(audio, {**start_options, **control})
    except WebSocketDisconnect:
        if active and not active.done():
            active.cancel()
        return