Newer
Older
navi-1 / clients / terminal / tui / chat_model.py
"""In-memory model for the chat panel.

Converts raw WebSocket events into logical chat messages that renderers understand.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Literal


@dataclass
class ChatItem:
    """Logical item shown in the chat history."""

    kind: Literal[
        "user_message",
        "assistant_message",
        "thinking_block",
        "tool_started",
        "tool_call",
        "error",
        "status",
    ]
    content: str = ""
    meta: dict = field(default_factory=dict)
    id: str = ""


class ChatModel:
    """Accumulate stream deltas and tool events into ChatItems."""

    def __init__(self) -> None:
        self.items: list[ChatItem] = []
        self._current_assistant: ChatItem | None = None
        self._current_thinking: ChatItem | None = None

    def add_user_message(self, text: str) -> None:
        self.items.append(ChatItem(kind="user_message", content=text))
        self._current_assistant = None

    def handle_ws_event(self, msg: dict) -> ChatItem | None:
        msg_type = msg.get("type")

        if msg_type == "stream_start":
            self._current_assistant = ChatItem(kind="assistant_message", content="")
            self.items.append(self._current_assistant)
            return None

        if msg_type == "thinking_delta":
            if self._current_thinking is None:
                self._current_thinking = ChatItem(kind="thinking_block", content="")
                self.items.append(self._current_thinking)
            self._current_thinking.content += msg.get("delta", "")
            return None

        if msg_type == "thinking_end":
            self._current_thinking = None
            return None

        if msg_type == "stream_delta":
            if self._current_assistant is None:
                self._current_assistant = ChatItem(kind="assistant_message", content="")
                self.items.append(self._current_assistant)
            self._current_assistant.content += msg.get("delta", "")
            return None

        if msg_type == "tool_started":
            item = ChatItem(kind="tool_started", meta=msg)
            self.items.append(item)
            return item

        if msg_type == "tool_call":
            item = ChatItem(kind="tool_call", meta=msg)
            self.items.append(item)
            return item

        if msg_type == "error":
            item = ChatItem(kind="error", content=msg.get("message", ""))
            self.items.append(item)
            return item

        if msg_type in ("stream_end", "context_compressed", "heartbeat", "session_sync"):
            return None

        # Unknown event — store as status for debugging.
        item = ChatItem(kind="status", content=f"{msg_type}: {msg}")
        self.items.append(item)
        return item