feat: update Dockerfile and pyproject.toml for dependency upgrades and improved websocket client

scripts/chit_chat_with_llm.py

@@ -1,175 +1,454 @@
-#!/usr/bin/env python
-"""Manual testing script for the websocket chat endpoint.
+#!/usr/bin/env python3
+"""Interactive Rich client for the FastAPI websocket chat endpoint.
 
-This script allows you to interact with the chat websocket endpoint manually.
-It connects to the endpoint and provides an interactive interface to:
-- Start conversations
-- Send messages
-- View full history
-- See streamed responses in real-time
+The script keeps the original websocket wire format intact while making the
+manual testing experience easier to discover and safer to operate.
 
-Usage:
-    python test_websocket_manual.py
+Examples:
+    python scripts/chit_chat_with_llm.py
+    python scripts/chit_chat_with_llm.py --url ws://localhost:8080/v1/chat/ws
+    python scripts/chit_chat_with_llm.py --debug --system-prompt "Be concise."
 
-Commands:
-    /start           - Start a new conversation (optionally with a system prompt)
-    /history         - Request the full message history
-    /end             - End the conversation
-    /help            - Show this help message
-    (any other text) - Send as a user message
+Interactive commands:
+    /start           Start or restart a conversation.
+    /history         Show the full message history.
+    /end, /quit      Gracefully close the websocket.
+    /help            Show the interactive help table.
+    any other text   Send text as a user message.
 """
 
+import argparse
 import asyncio
 import json
-import sys
-from typing import Any
+import logging
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import Any, Literal, Protocol, TypedDict, cast
 
 import websockets
-from rich import print as rprint
-from rich.console import Console
+from rich import box
+from rich.console import Console, Group
+from rich.logging import RichHandler
+from rich.panel import Panel
+from rich.progress import Progress, SpinnerColumn, TextColumn, TimeElapsedColumn
 from rich.table import Table
+from rich.text import Text
+from rich.traceback import install as install_rich_traceback
+from websockets.exceptions import ConnectionClosed, InvalidURI, WebSocketException
 
-console = Console()
+DEFAULT_WS_URL = "ws://localhost:8080/v1/chat/ws"
+DEFAULT_SYSTEM_PROMPT = "You are a helpful assistant."
 
-# Configuration
-WS_URL = "ws://localhost:8080/v1/chat/ws"
+EventType = Literal[
+    "connected",
+    "started",
+    "assistant_chunk",
+    "assistant_message",
+    "history",
+    "error",
+]
+ClientMessageType = Literal["start", "user_message", "history_request", "end"]
 
 
-async def send_message(websocket: websockets.WebSocketClientProtocol, message: dict[str, Any]) -> None:
-    """Send a message to the websocket."""
-    await websocket.send(json.dumps(message))
+class WebSocketLike(Protocol):
+    """Minimal websocket protocol used by this client."""
+
+    async def send(self, message: str) -> None: ...
+
+    async def recv(self) -> str | bytes: ...
 
 
-async def receive_and_display_event(websocket: websockets.WebSocketClientProtocol) -> dict[str, Any]:
-    """Receive an event from websocket and display it."""
-    data = await websocket.recv()
-    event = json.loads(data)
+class ChatMessage(TypedDict):
+    """Role/content message shape returned by the websocket API."""
 
-    if event["type"] == "connected":
-        rprint(f"[green]✓ Connected[/green] - Session ID: {event['session_id']}")
-    elif event["type"] == "started":
-        rprint(f"[green]✓ Conversation started[/green] - Session ID: {event['session_id']}")
-    elif event["type"] == "assistant_chunk":
-        # Print chunk without newline for streaming effect
-        console.print(event["content"], end="", soft_wrap=True)
-    elif event["type"] == "assistant_message":
-        console.print()  # New line after chunks
-        rprint(f"[blue]✓ Message complete[/blue]")
-    elif event["type"] == "history":
-        rprint("\n[cyan]Message History:[/cyan]")
-        table = Table(show_header=True, header_style="bold magenta")
-        table.add_column("Role", style="cyan", width=12)
-        table.add_column("Content", style="green")
-
-        for msg in event["messages"]:
-            role = msg["role"]
-            content = msg["content"][:100] + "..." if len(msg["content"]) > 100 else msg["content"]
-            table.add_row(role, content)
-
-        console.print(table)
-    elif event["type"] == "error":
-        rprint(f"[red]✗ Error: {event.get('message', 'Unknown error')}[/red]")
-    else:
-        rprint(f"[yellow]? Received event: {event}[/yellow]")
-
-    return event
+    role: Literal["system", "user", "assistant"]
+    content: str
 
 
-async def get_user_input(prompt: str = "") -> str:
-    """Get user input asynchronously."""
-    loop = asyncio.get_event_loop()
-    return await loop.run_in_executor(None, input, prompt)
+class ServerEvent(TypedDict, total=False):
+    """Known server event fields.
+
+    The websocket endpoint uses a discriminated JSON shape. ``total=False``
+    keeps display code resilient to malformed or future server events.
+    """
+
+    type: EventType
+    session_id: str
+    index: int
+    content: str
+    message: ChatMessage
+    messages: list[ChatMessage]
 
 
-async def main():
-    """Main function for interactive websocket testing."""
-    rprint("[bold cyan]FastAPI WebSocket Chat Tester[/bold cyan]")
-    rprint(f"Connecting to {WS_URL}...\n")
+@dataclass(frozen=True, slots=True)
+class ChatClientConfig:
+    """Runtime settings parsed from the command line."""
 
-    try:
-        async with websockets.connect(WS_URL) as websocket:
-            # Receive the initial connected message
-            connected_event = await receive_and_display_event(websocket)
-            session_id = connected_event["session_id"]
+    url: str = DEFAULT_WS_URL
+    system_prompt: str = DEFAULT_SYSTEM_PROMPT
+    connect_timeout: float = 10.0
+    open_on_connect: bool = False
+    debug: bool = False
 
-            rprint("\n[yellow]Commands: /start, /history, /end, /help, or type a message[/yellow]\n")
 
-            while True:
-                try:
-                    # Get user input
-                    user_input = await get_user_input("[bold]You:[/bold] ")
+@dataclass(slots=True)
+class ChatClientState:
+    """Mutable client state for the active websocket session."""
 
-                    if user_input.strip() == "":
-                        continue
+    session_id: str | None = None
+    conversation_started: bool = False
+    received_chunks: int = 0
 
-                    # Handle special commands
-                    if user_input.lower() == "/help":
-                        rprint("\n[cyan]Available Commands:[/cyan]")
-                        rprint("[green]/start[/green]    - Start a new conversation")
-                        rprint("[green]/history[/green]  - Request the full message history")
-                        rprint("[green]/end[/green]      - End the conversation")
-                        rprint("[green]/help[/green]     - Show this help message\n")
-                        continue
 
-                    elif user_input.lower() == "/start":
-                        system_prompt = await get_user_input(
-                            "Enter system prompt (or press Enter for default): "
-                        )
-                        msg = {
-                            "type": "start",
-                            "system_prompt": system_prompt or "You are a helpful assistant.",
-                        }
-                        await send_message(websocket, msg)
-                        event = await receive_and_display_event(websocket)
-                        rprint()
-                        continue
+class ChatCli:
+    """UX-focused websocket chat client."""
 
-                    elif user_input.lower() == "/history":
-                        msg = {"type": "history_request"}
-                        await send_message(websocket, msg)
-                        event = await receive_and_display_event(websocket)
-                        rprint()
-                        continue
+    def __init__(self, config: ChatClientConfig, console: Console) -> None:
+        self.config = config
+        self.console = console
+        self.state = ChatClientState()
 
-                    elif user_input.lower() == "/end":
-                        msg = {"type": "end"}
-                        await send_message(websocket, msg)
-                        rprint("[yellow]Conversation ended. Closing connection...[/yellow]")
-                        break
+    async def run(self) -> int:
+        """Connect to the websocket and start the interactive command loop."""
+        self._show_welcome()
 
-                    else:
-                        # Send as a regular user message
-                        msg = {"type": "user_message", "content": user_input}
-                        await send_message(websocket, msg)
-                        rprint("[cyan]Assistant:[/cyan] ", end="")
+        try:
+            with self._progress(f"Connecting to [bold]{self.config.url}[/bold]..."):
+                websocket = await asyncio.wait_for(
+                    websockets.connect(self.config.url),
+                    timeout=self.config.connect_timeout,
+                )
 
-                        # Receive chunks and the final message
-                        while True:
-                            event = await receive_and_display_event(websocket)
-                            if event["type"] == "assistant_message":
-                                break
+            async with websocket:
+                connected_event = await self.receive_event(websocket)
+                self.display_event(connected_event)
 
-                        rprint()  # New line after response
+                if self.config.open_on_connect:
+                    await self.start_conversation(websocket, self.config.system_prompt)
 
-                except KeyboardInterrupt:
-                    rprint("\n[yellow]Interrupted by user.[/yellow]")
-                    msg = {"type": "end"}
-                    await send_message(websocket, msg)
+                await self.interactive_loop(websocket)
+                return 0
+        except TimeoutError:
+            self.print_error(
+                "Connection timed out.",
+                f"No websocket accepted the connection within {self.config.connect_timeout:g} seconds. "
+                "Check that the FastAPI server is running and reachable.",
+            )
+        except ConnectionRefusedError:
+            self.print_error(
+                "Connection refused.",
+                "Start the API server first, then try again. Expected default: http://localhost:8080",
+            )
+        except InvalidURI as exc:
+            self.print_error(
+                "Invalid websocket URL.",
+                f"{exc}. Use a URL such as ws://localhost:8080/v1/chat/ws.",
+            )
+        except ConnectionClosed as exc:
+            close_code = exc.rcvd.code if exc.rcvd else exc.sent.code if exc.sent else "unknown"
+            close_reason = exc.rcvd.reason if exc.rcvd else exc.sent.reason if exc.sent else "none"
+            self.print_warning(
+                f"The websocket closed unexpectedly (code={close_code}, reason={close_reason or 'none'})."
+            )
+        except KeyboardInterrupt:
+            self.print_warning("Interrupted by user.")
+        except WebSocketException as exc:
+            self.print_error(
+                "Websocket error.",
+                f"{exc}. Verify the endpoint path and server logs for details.",
+            )
+
+        return 1
+
+    async def interactive_loop(self, websocket: WebSocketLike) -> None:
+        """Read commands from the terminal until the user exits."""
+        self.print_info("Type [bold cyan]/help[/bold cyan] to see available commands.")
+
+        while True:
+            user_input = (await self.get_user_input("[bold green]You[/bold green] › ")).strip()
+
+            if not user_input:
+                continue
+
+            command = user_input.casefold()
+            match command:
+                case "/help":
+                    self.show_command_help()
+                case "/start":
+                    system_prompt = await self.get_user_input(
+                        "[bold yellow]System prompt[/bold yellow] (Enter for default) › "
+                    )
+                    await self.start_conversation(websocket, system_prompt or self.config.system_prompt)
+                case "/history":
+                    await self.request_history(websocket)
+                case "/end" | "/quit" | "/exit":
+                    await self.end_conversation(websocket)
                     break
+                case _:
+                    await self.send_user_message(websocket, user_input)
 
-    except ConnectionRefusedError:
-        rprint(
-            f"[red]✗ Could not connect to {WS_URL}[/red]"
+    async def start_conversation(self, websocket: WebSocketLike, system_prompt: str) -> None:
+        """Start or restart the server-side conversation history."""
+        await self.send_message(websocket, {"type": "start", "system_prompt": system_prompt})
+        with self._progress("Waiting for conversation acknowledgement..."):
+            event = await self.receive_event(websocket)
+        self.display_event(event)
+        self.console.print()
+
+    async def request_history(self, websocket: WebSocketLike) -> None:
+        """Ask the server for complete message history."""
+        await self.send_message(websocket, {"type": "history_request"})
+        with self._progress("Loading message history..."):
+            event = await self.receive_event(websocket)
+        self.display_event(event)
+        self.console.print()
+
+    async def end_conversation(self, websocket: WebSocketLike) -> None:
+        """Gracefully close the server-side conversation."""
+        await self.send_message(websocket, {"type": "end"})
+        self.print_warning("Conversation ended. Closing connection...")
+
+    async def send_user_message(self, websocket: WebSocketLike, content: str) -> None:
+        """Send one user message and stream the assistant response."""
+        await self.send_message(websocket, {"type": "user_message", "content": content})
+        self.console.print("[bold cyan]Assistant[/bold cyan] › ", end="")
+
+        self.state.received_chunks = 0
+        while True:
+            event = await self.receive_event(websocket)
+            event_type = event.get("type")
+            self.display_event(event)
+
+            if event_type == "assistant_message":
+                break
+            if event_type == "error":
+                self.console.print()
+                break
+            if event_type != "assistant_chunk":
+                self.print_warning(
+                    f"Expected assistant_chunk or assistant_message, received {event_type!r}."
+                )
+                break
+
+        self.console.print()
+
+    async def send_message(
+        self,
+        websocket: WebSocketLike,
+        message: Mapping[str, Any],
+    ) -> None:
+        """Serialize and send one client event."""
+        await websocket.send(json.dumps(message))
+        logging.getLogger(__name__).debug("Sent websocket event: %s", message)
+
+    async def receive_event(self, websocket: WebSocketLike) -> ServerEvent:
+        """Receive and decode one JSON server event."""
+        raw_event = await websocket.recv()
+        if isinstance(raw_event, bytes):
+            raw_event = raw_event.decode()
+
+        try:
+            decoded = json.loads(raw_event)
+        except json.JSONDecodeError as exc:
+            raise RuntimeError(f"Server returned invalid JSON: {raw_event!r}") from exc
+
+        if not isinstance(decoded, dict):
+            raise RuntimeError(f"Server returned a non-object event: {decoded!r}")
+
+        logging.getLogger(__name__).debug("Received websocket event: %s", decoded)
+        return cast(ServerEvent, cast(object, decoded))
+
+    def display_event(self, event: ServerEvent) -> None:
+        """Render one server event with semantic styling."""
+        match event.get("type"):
+            case "connected":
+                self.state.session_id = event.get("session_id")
+                self.print_success(f"Connected. Session ID: [bold]{self.state.session_id}[/bold]")
+            case "started":
+                self.state.conversation_started = True
+                self.print_success(f"Conversation started. Session ID: [bold]{event.get('session_id')}[/bold]")
+            case "assistant_chunk":
+                self.state.received_chunks += 1
+                self.console.print(event.get("content", ""), end="", soft_wrap=True)
+            case "assistant_message":
+                self.console.print()
+                message = event.get("message", {})
+                content = message.get("content", "") if isinstance(message, dict) else ""
+                self.print_success(
+                    f"Message complete ({self.state.received_chunks} chunk(s), {len(content)} character(s))."
+                )
+            case "history":
+                self.render_history(event.get("messages", []))
+            case "error":
+                self.print_error("Server returned an error.", str(event.get("message", "Unknown error")))
+            case unknown:
+                self.console.print(
+                    Panel.fit(
+                        json.dumps(event, indent=2, default=str),
+                        title=f"Unknown event: {unknown!r}",
+                        border_style="yellow",
+                    )
+                )
+
+    def render_history(self, messages: list[ChatMessage] | object) -> None:
+        """Render conversation history as a Rich table."""
+        if not isinstance(messages, list):
+            self.print_warning("History response did not include a valid messages list.")
+            return
+
+        table = Table(
+            title=f"Message History ({len(messages)} message(s))",
+            box=box.ROUNDED,
+            header_style="bold magenta",
+            show_lines=True,
         )
-        rprint("[yellow]Make sure the FastAPI server is running on http://localhost:8080[/yellow]")
-        sys.exit(1)
-    except Exception as e:
-        rprint(f"[red]✗ Error: {e}[/red]")
-        sys.exit(1)
-    finally:
-        rprint("[cyan]Connection closed.[/cyan]")
+        table.add_column("#", justify="right", style="dim", width=4)
+        table.add_column("Role", style="cyan", width=12)
+        table.add_column("Content", style="green", overflow="fold")
+
+        for index, message in enumerate(messages, start=1):
+            role = str(message.get("role", "unknown")) if isinstance(message, dict) else "unknown"
+            content = str(message.get("content", "")) if isinstance(message, dict) else repr(message)
+            table.add_row(str(index), role, content)
+
+        self.console.print(table)
+
+    async def get_user_input(self, prompt: str) -> str:
+        """Read terminal input without blocking the event loop."""
+        return await asyncio.to_thread(self.console.input, prompt)
+
+    def show_command_help(self) -> None:
+        """Display interactive command help."""
+        table = Table(title="Interactive Commands", box=box.SIMPLE_HEAVY)
+        table.add_column("Command", style="bold cyan", no_wrap=True)
+        table.add_column("Action", style="white")
+        table.add_row("/start", "Start or restart a conversation; clears server-side history.")
+        table.add_row("/history", "Request and display the full role/content message history.")
+        table.add_row("/end, /quit, /exit", "Gracefully close the websocket connection.")
+        table.add_row("/help", "Show this help table.")
+        table.add_row("any text", "Send text as a user message and stream the reply.")
+        self.console.print(table)
+
+    def _show_welcome(self) -> None:
+        command_summary = Table.grid(padding=(0, 2))
+        command_summary.add_column(style="bold cyan", no_wrap=True)
+        command_summary.add_column(style="white")
+        command_summary.add_row("Endpoint", self.config.url)
+        command_summary.add_row("Default prompt", self.config.system_prompt)
+        command_summary.add_row("Start mode", "auto" if self.config.open_on_connect else "manual (/start)")
+
+        self.console.print(
+            Panel(
+                Group(
+                    Text("FastAPI WebSocket Chat Tester", style="bold cyan"),
+                    Text("Stream replies, inspect history, and validate the chat websocket UX."),
+                    command_summary,
+                ),
+                border_style="cyan",
+                padding=(1, 2),
+            )
+        )
+
+    def _progress(self, message: str) -> Progress:
+        """Create a compact indeterminate Rich progress spinner."""
+        progress = Progress(
+            SpinnerColumn(style="cyan"),
+            TextColumn(message),
+            TimeElapsedColumn(),
+            console=self.console,
+            transient=True,
+        )
+        progress.add_task("operation", total=None)
+        return progress
+
+    def print_info(self, message: str) -> None:
+        self.console.print(f"[bold blue]ℹ[/bold blue] {message}")
+
+    def print_success(self, message: str) -> None:
+        self.console.print(f"[bold green]✓[/bold green] {message}")
+
+    def print_warning(self, message: str) -> None:
+        self.console.print(f"[bold yellow]⚠[/bold yellow] {message}")
+
+    def print_error(self, title: str, detail: str) -> None:
+        self.console.print(
+            Panel(
+                detail,
+                title=f"[bold red]✗ {title}[/bold red]",
+                border_style="red",
+                expand=False,
+            )
+        )
+
+
+def parse_args(argv: list[str] | None = None) -> ChatClientConfig:
+    """Parse CLI options into a typed configuration object."""
+    parser = argparse.ArgumentParser(
+        description="Interactive Rich client for the FastAPI websocket chat endpoint.",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument(
+        "--url",
+        default=DEFAULT_WS_URL,
+        help="Websocket endpoint to connect to.",
+    )
+    parser.add_argument(
+        "--system-prompt",
+        default=DEFAULT_SYSTEM_PROMPT,
+        help="Default system prompt used by /start or --start.",
+    )
+    parser.add_argument(
+        "--connect-timeout",
+        type=float,
+        default=10.0,
+        help="Seconds to wait while opening the websocket connection.",
+    )
+    parser.add_argument(
+        "--start",
+        action="store_true",
+        dest="open_on_connect",
+        help="Automatically send a start event after connecting.",
+    )
+    parser.add_argument(
+        "--debug",
+        action="store_true",
+        help="Enable Rich-formatted debug logging and local tracebacks.",
+    )
+
+    namespace = parser.parse_args(argv)
+    return ChatClientConfig(
+        url=namespace.url,
+        system_prompt=namespace.system_prompt,
+        connect_timeout=namespace.connect_timeout,
+        open_on_connect=namespace.open_on_connect,
+        debug=namespace.debug,
+    )
+
+
+def configure_logging(debug: bool) -> None:
+    """Install Rich tracebacks and configure optional debug logging."""
+    install_rich_traceback(show_locals=debug, suppress=[websockets])
+    logging.basicConfig(
+        level=logging.DEBUG if debug else logging.WARNING,
+        format="%(message)s",
+        datefmt="[%X]",
+        handlers=[RichHandler(rich_tracebacks=True, show_path=debug)],
+    )
+
+
+async def async_main(argv: list[str] | None = None) -> int:
+    """Async application entry point."""
+    config = parse_args(argv)
+    configure_logging(config.debug)
+    console = Console()
+    return await ChatCli(config=config, console=console).run()
+
+
+def main() -> None:
+    """Synchronous script entry point."""
+    raise SystemExit(asyncio.run(async_main()))
 
 
 if __name__ == "__main__":
-    asyncio.run(main())
+    main()