my_voice_assistant_v2/app/api/ws.py

"""WebSocket-Streaming-Chat (Echtzeit-Transport, erster Increment).

Etabliert einen dauerhaften, bidirektionalen Kanal: der Client schickt pro Turn
eine JSON-Nachricht, der Server streamt strukturierte Events zurueck
(ack -> semantic -> audio (binaer) -> done). Auth, Session-Gedaechtnis und
Langzeit-Erinnerungen gelten wie bei POST /api/chat.

Bewusst spaeter (eigene Increments): Token-Level-LLM-Streaming, Audio-Eingang/
Streaming-STT, Barge-in/Interrupt und WebRTC.
"""

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()


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)


@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:
        # Vor accept() schliessen -> Handshake wird mit 403 abgelehnt.
        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

            overrides = {
                key: msg.get(key)
                for key in (
                    "input_endpoint",
                    "output_endpoint",
                    "language",
                    "stt_provider",
                    "llm_provider",
                    "tts_provider",
                )
            }

            try:
                route = resolve_route(user, session_id, overrides)
                orchestrator = build_orchestrator(route)
                output = await resolve_output_endpoint(route)
                conversation = (
                    store.get_recent_messages(session_id, settings.history_max_messages)
                    if session_id
                    else []
                )
                memories = store.get_memories(user.id)
            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

            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 = msg.get("voice") or settings.openrouter_tts_voice
            stream = bool(msg.get("stream"))
            audio_stream = bool(msg.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)})
                continue

            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,
                }
            )
            # Bei audio_stream wurden die Audio-Chunks bereits live gesendet.
            if not audio_stream:
                await websocket.send_bytes(audio)
            await websocket.send_json(
                {"type": "done", "audio_format": "pcm", "sample_rate": 24000}
            )
    except WebSocketDisconnect:
        return