src/session_config.py

# SPDX-FileCopyrightText: Copyright (c) 2026 Abhinav Kalvacherla
# SPDX-License-Identifier: Apache-2.0
"""
Per-session config slot for MoshiRAG dynamic function-calling.

A client supplies, at WebSocket connect time, a prompt + a list of declarative
function/API definitions (see http_executor.py for the per-function schema).
The launcher (model.py) captures that config from the connect header in a
monkey-patched `ServerState.handle_chat` and stores it here; the singleton
`APIToolReferenceGenerator` reads it on each RAG trigger.

Concurrency note: the deployment runs `predict_concurrency: 1` (one Moshi
conversation per replica), so a single global slot is safe — there is never
more than one active connection. To support concurrency > 1, replace the slot
with a dict keyed by a session id and thread that id through the per-channel
`RAGManager._active_profile_id` (set via `set_retrieval_profile_id`), which is
already plumbed into `generate_reference_text(..., active_profile_id=...)`.
"""

from __future__ import annotations

import logging
import threading
from dataclasses import dataclass, field

logger = logging.getLogger(__name__)

_VALID_METHODS = {"GET", "POST", "PUT", "PATCH", "DELETE"}


@dataclass
class SessionConfig:
    prompt: str
    functions: list[dict] = field(default_factory=list)
    llm_base_url: str | None = None
    llm_model_name: str | None = None
    allowed_hosts: list[str] | None = None
    filler: str | None = None   # spoken while an API call is in flight


_lock = threading.Lock()
_current: SessionConfig | None = None


def set_session_config(raw: dict | None) -> None:
    """Validate and store the config for the current connection (or clear it)."""
    global _current
    cfg = _validate(raw) if raw else None
    with _lock:
        _current = cfg
    if cfg is not None:
        logger.info(
            "[session] config set: prompt_len=%d functions=%s",
            len(cfg.prompt), [f["name"] for f in cfg.functions],
        )


def clear_session_config() -> None:
    global _current
    with _lock:
        _current = None
    logger.info("[session] config cleared")


def get_session_config() -> SessionConfig | None:
    with _lock:
        return _current


def _valid_func(f: object) -> bool:
    return (
        isinstance(f, dict)
        and isinstance(f.get("name"), str)
        and f["name"].strip() != ""
        and isinstance(f.get("endpoint"), str)
        and str(f.get("method", "GET")).upper() in _VALID_METHODS
    )


def _validate(raw: dict) -> SessionConfig:
    prompt = str(raw.get("prompt", "")).strip()
    funcs_in = raw.get("functions") or []
    clean: list[dict] = []
    for f in funcs_in:
        if _valid_func(f):
            clean.append(f)
        else:
            logger.warning("[session] dropping malformed function def: %r", f)
    return SessionConfig(
        prompt=prompt,
        functions=clean,
        llm_base_url=raw.get("llm_base_url"),
        llm_model_name=raw.get("llm_model_name"),
        allowed_hosts=raw.get("allowed_hosts"),
        filler=raw.get("filler"),
    )