state.py

from typing import Literal, NotRequired, TypedDict

import streamlit as st

from utils.helpers import session_key

PendingChatAction = Literal["new_user_prompt", "regenerate_after_edit"]


class ChatMessage(TypedDict):
    role: str
    content: str
    _contrast: NotRequired[object]
    _needs_contrast: NotRequired[bool]


class ChatState(TypedDict):
    messages: list[ChatMessage]
    persona_id: str | None
    prompt_mode: str


def chat_session_key(model_name: str, dataset_source: str) -> str:
    """Build the session-state key for a chat conversation.

    A model/backend switch changes *how* the next turn is generated, not which
    conversation the user is looking at. Keeping the model out of the key means
    toggling local/remote execution (or selecting another model) no longer makes
    an existing thread appear to vanish behind a fresh empty state.

    ``model_name`` stays in the signature for call-site compatibility and to
    make the intent explicit where chat state is requested.
    """

    _ = model_name
    return session_key("chat_state", dataset_source)


def default_chat_state() -> ChatState:
    return {
        "messages": [],
        "persona_id": None,
        "prompt_mode": "templated",
    }


def reset_chat_context_state(
    state: ChatState,
    persona_id: str,
    prompt_mode: str,
    *ui_keys: str,
) -> None:
    """Reset one chat context and clear any related widget state."""

    state["messages"] = []
    state["persona_id"] = persona_id
    state["prompt_mode"] = prompt_mode
    for key in ui_keys:
        st.session_state.pop(key, None)


def get_chat_state(model_name: str, dataset_source: str) -> ChatState:
    """Return the mutable chat state for the active context."""

    key = chat_session_key(model_name, dataset_source)
    return st.session_state.setdefault(key, default_chat_state())