Source code for chariot.cli.awm.chat.menu

"""Chat menu entrypoint for AWM CLI."""

import click
import questionary

from chariot import awm
from chariot.cli.awm.chat.loop import _run_chat_loop
from chariot.cli.awm.chat.render import _render_session_history
from chariot.cli.awm.chat.sessions import (
    _fetch_session_history,
    _list_user_sessions,
    _pick_agent,
    _session_column_widths,
    _session_row_title,
)
from chariot.cli.awm.common import (
    _BACK,
    _NEW,
    _QUESTIONARY_STYLE,
    _flush_stdin,
    _run_menu_action,
)
from chariot.cli.awm.status import _get_workflow_paused_state

_SESSION_HEADER_WITH_AGENTS = "__awm_session_header_with_agents__"




def do_chat_from_menu(workflow: awm.models.Workflow) -> None:
    """Open the interactive chat session picker for a workflow."""
    paused_state = _run_menu_action(
        lambda: _get_workflow_paused_state(workflow.id),
        action_label=f"get paused state for workflow '{workflow.name}'",
        fallback=None,
    )
    if paused_state is True:
        click.echo(
            click.style("  Workflow is paused. ", fg="yellow")
            + "Resume it first before starting a chat session."
        )
        return

    sessions = _run_menu_action(
        lambda: _list_user_sessions(workflow.id),
        action_label=f"list sessions for workflow '{workflow.name}'",
        fallback=[],
    )

    choices: list[questionary.Choice] = [
        questionary.Choice(title="New session", value=_NEW),
    ]
    if sessions:
        time_w, id_w, count_w = _session_column_widths(sessions)
        choices.append(questionary.Separator("─── Previous sessions ───"))
        choices.append(
            questionary.Choice(
                title=(
                    f"{'Session Time'.ljust(time_w)}  "
                    f"{'Session ID'.ljust(id_w)}  "
                    f"{'Msgs'.rjust(count_w)}  Agents"
                ),
                value=_SESSION_HEADER_WITH_AGENTS,
            )
        )
        choices.append(questionary.Separator())
        for session in sessions:
            choices.append(
                questionary.Choice(
                    title=_session_row_title(
                        session,
                        time_w=time_w,
                        id_w=id_w,
                        count_w=count_w,
                        include_agents=True,
                    ),
                    value=session,
                )
            )
    choices.append(questionary.Separator())
    choices.append(questionary.Choice(title="↩ Cancel", value=_BACK))

    _flush_stdin()
    while True:
        selected = questionary.select(
            "Select session:",
            choices=choices,
            style=_QUESTIONARY_STYLE,
        ).ask()
        if selected != _SESSION_HEADER_WITH_AGENTS:
            break

    if selected is None or selected == _BACK:
        return

    if selected == _NEW:
        agent_name = _run_menu_action(
            lambda: _pick_agent(workflow),
            action_label=f"select an agent for workflow '{workflow.name}'",
            fallback=None,
        )
        if not agent_name:
            return
        _run_chat_loop(workflow, agent_name, None)
        return

    # Resuming an existing session — derive agent from session metadata.
    session: awm.models.SessionDetails = selected
    if len(session.agents) == 1:
        agent_name = session.agents[0]
    elif len(session.agents) > 1:
        _flush_stdin()
        agent_name = questionary.select(
            "Select agent for this session:",
            choices=[questionary.Choice(title=name, value=name) for name in session.agents]
            + [questionary.Choice(title="↩ Cancel", value=_BACK)],
            style=_QUESTIONARY_STYLE,
        ).ask()
        if agent_name is None or agent_name == _BACK:
            return
    else:
        agent_name = _run_menu_action(
            lambda: _pick_agent(workflow),
            action_label=f"select an agent for workflow '{workflow.name}'",
            fallback=None,
        )
        if not agent_name:
            return

    history = _run_menu_action(
        lambda: _fetch_session_history(workflow.id, session.session_id),
        action_label=f"load session history for workflow '{workflow.name}'",
        fallback=None,
    )
    if history is None:
        return
    _render_session_history(history, agent_name)
    _run_chat_loop(workflow, agent_name, session.session_id)