my_voice_assistant_v3/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.

Event-Folge der Antwort: 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.

Spaeter (eigene Increments): partielle Live-Transkripte (Streaming-STT mit VAD),
Barge-in/Turn-Manager und WebRTC.
"""

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,
)
from app.store import SessionOwnershipError

router = APIRouter()

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


def _authenticate(token: str | None):
    store = get_store()
    if not settings.auth_enabled:
        return store.ensure_anonymous_user()
    if not token:
        return None
    return store.get_user_by_token(token)


async def _resolve(user, session_id, options):
    """Loest Route + Orchestrator + Output-Endpunkt auf (kann RoutingError/Ownership werfen)."""
    overrides = {key: options.get(key) for key in _OVERRIDE_KEYS}
    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):
    """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)
    if memories:
        memory_text = "Was du ueber den Nutzer weisst:\n" + "\n".join(
            f"- {m.content}" for m in memories
        )
        llm_context = [{"role": "system", "content": memory_text}] + llm_context

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

    voice = options.get("voice") or settings.openrouter_tts_voice
    stream = bool(options.get("stream"))
    audio_stream = bool(options.get("audio_stream"))

    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)

    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,
            )
        else:
            trace, audio = await orchestrator.chat_text(
                text,
                language=route.language,
                voice=voice,
                output=output,
                history=llm_context,
            )
    except Exception as exc:
        await websocket.send_json({"type": "error", "status": 502, "detail": str(exc)})
        return

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

    await websocket.send_json(
        {"type": "semantic", "text": trace.semantic_response, "spoken": trace.spoken_response}
    )
    if not audio_stream:
        await websocket.send_bytes(audio)
    await websocket.send_json({"type": "done", "audio_format": "pcm", "sample_rate": 24000})


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

    try:
        while True:
            msg = await websocket.receive_json()
            text = (msg.get("text") or "").strip()
            if not text:
                await websocket.send_json({"type": "error", "detail": "empty text"})
                continue
            try:
                route, orchestrator, output = await _resolve(user, session_id, msg)
            except SessionOwnershipError as exc:
                await websocket.send_json({"type": "error", "status": 403, "detail": str(exc)})
                continue
            except RoutingError as exc:
                await websocket.send_json({"type": "error", "status": 422, "detail": str(exc)})
                continue
            await _run_turn(websocket, store, user, session_id, route, orchestrator, output, text, msg)
    except WebSocketDisconnect:
        return


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

    audio_buffer = bytearray()
    fmt = "wav"

    try:
        while True:
            message = await websocket.receive()
            if message["type"] == "websocket.disconnect":
                return

            if message.get("bytes") is not None:
                audio_buffer.extend(message["bytes"])
                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 == "start":
                audio_buffer.clear()
                fmt = control.get("format", "wav")
                continue
            if ctype != "end":
                continue

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

            try:
                route, orchestrator, output = await _resolve(user, session_id, control)
            except SessionOwnershipError as exc:
                await websocket.send_json({"type": "error", "status": 403, "detail": str(exc)})
                audio_buffer.clear()
                continue
            except RoutingError as exc:
                await websocket.send_json({"type": "error", "status": 422, "detail": str(exc)})
                audio_buffer.clear()
                continue

            try:
                transcript = await orchestrator.stt.transcribe(
                    bytes(audio_buffer), fmt=fmt, language=route.language
                )
            except Exception as exc:
                await websocket.send_json({"type": "error", "status": 502, "detail": str(exc)})
                audio_buffer.clear()
                continue
            finally:
                audio_buffer.clear()

            await websocket.send_json({"type": "transcript", "text": transcript})
            if not transcript or not transcript.strip():
                await websocket.send_json({"type": "error", "detail": "empty transcript"})
                continue

            await _run_turn(
                websocket, store, user, session_id, route, orchestrator, output, transcript, control
            )
    except WebSocketDisconnect:
        return