prain: terminal UI

prain/tui.py

@@ -0,0 +1,115 @@
+"""prain tui: a terminal chat that warms experts up while you type. Stdlib curses."""
+
+import curses
+import locale
+import os
+import queue
+import textwrap
+import threading
+import time
+
+from prain.chat import (ABORTED, Warmer, chat_body, first_model, post, send_request,
+                          stream_text, warm_body)
+
+WARMUP_MODES = ("idle", "always", "off")
+
+
+class Chat:
+    """Conversation state without curses, so tests can drive it.
+    warmup: "idle" warms up only while no answer streams, "always" also during one, "off" never."""
+
+    def __init__(self, base_url, model, warmup="idle", priority=1):
+        self.base_url, self.model, self.warmup = base_url, model, warmup
+        self.messages = []  # finished turns
+        self.draft = ""
+        self.reply = None  # the answer text while it streams, else None
+        self.status = "type a message; Enter sends, Esc quits"
+        self.events = queue.Queue()
+        self.warmer = Warmer(lambda draft: send_request(
+            base_url, "/v1/chat/completions", warm_body(model, self.messages, draft, priority)))
+
+    def type(self, text):
+        self.draft = text
+        idle = self.reply is None
+        if text.strip() and (self.warmup == "always" or (self.warmup == "idle" and idle)):
+            self.warmer.update(text)
+
+    def send(self):
+        if self.reply is not None or not self.draft.strip():
+            return
+        self.warmer.cancel()  # a queued draft is stale now; the in-flight one only delays us
+        self.messages.append({"role": "user", "content": self.draft.strip()})
+        self.draft, self.reply = "", ""
+        threading.Thread(target=self._stream, args=(list(self.messages),), daemon=True).start()
+
+    def _stream(self, messages):
+        start, first, pieces = time.monotonic(), None, 0
+        try:
+            with post(self.base_url, "/v1/chat/completions", chat_body(self.model, messages)) as r:
+                for piece in stream_text(r):
+                    first = first or time.monotonic()
+                    pieces += 1
+                    self.events.put(("piece", piece))
+            rate = (pieces - 1) / (time.monotonic() - first) if pieces > 1 else 0.0
+            ttft = (first or time.monotonic()) - start
+            self.events.put(("done", f"first token {ttft:.1f} s, {rate:.1f} tok/s"))
+        except ABORTED as e:
+            self.events.put(("done", f"error: {e}"))
+
+    def poll(self):
+        """Apply streamed pieces; call from the UI thread."""
+        while True:
+            try:
+                kind, value = self.events.get_nowait()
+            except queue.Empty:
+                return
+            if kind == "piece":
+                self.reply += value
+            else:
+                self.messages.append({"role": "assistant", "content": self.reply})
+                self.reply, self.status = None, value
+
+
+def draw(screen, chat):
+    screen.erase()
+    height, width = screen.getmaxyx()
+    turns = chat.messages + ([{"role": "assistant", "content": chat.reply}]
+                             if chat.reply is not None else [])
+    lines = []
+    for turn in turns:
+        prefix = "you> " if turn["role"] == "user" else "ai>  "
+        for paragraph in (prefix + turn["content"]).split("\n"):
+            lines += textwrap.wrap(paragraph, width - 1) or [""]
+    for y, line in enumerate(lines[-(height - 2):]):
+        screen.addstr(y, 0, line)
+    warming = "warming experts... " if chat.warmer.busy else ""
+    screen.addstr(height - 2, 0, f"[{chat.model}] {warming}{chat.status}"[: width - 1],
+                  curses.A_REVERSE)
+    screen.addstr(height - 1, 0, ("> " + chat.draft)[-(width - 1):])
+    screen.refresh()
+
+
+def run(screen, chat):
+    screen.timeout(50)  # redraw at least every 50 ms while an answer streams
+    while True:
+        chat.poll()
+        draw(screen, chat)
+        try:
+            key = screen.get_wch()
+        except curses.error:
+            continue  # nothing typed
+        if key in ("\x1b", "\x04"):  # Esc or Ctrl-D
+            return
+        if key in ("\n", "\r", curses.KEY_ENTER):
+            chat.send()
+        elif key in (curses.KEY_BACKSPACE, "\x7f", "\b"):
+            chat.type(chat.draft[:-1])
+        elif isinstance(key, str) and key.isprintable():
+            chat.type(chat.draft + key)
+
+
+def main(base_url, warmup="idle", priority=1):
+    locale.setlocale(locale.LC_ALL, "")
+    os.environ.setdefault("ESCDELAY", "25")  # Esc quits at once, not after 1 s
+    chat = Chat(base_url, first_model(base_url), warmup, priority)
+    curses.wrapper(run, chat)

tests/test_tui.py

@@ -0,0 +1,70 @@
+import fcntl
+import os
+import pty
+import struct
+import subprocess
+import sys
+import termios
+import threading
+import unittest
+
+from prain.tui import Chat
+from tests.helpers import StubVllm, wait_for
+
+
+class Tui(unittest.TestCase):
+    def setUp(self):
+        self.stub = StubVllm()
+        self.addCleanup(self.stub.close)
+
+    def test_typing_warms_up_and_enter_streams_the_answer(self):
+        chat = Chat(self.stub.url, "stub-model")
+        chat.type("Why is the sky")
+        wait_for(lambda: self.stub.bodies and not chat.warmer.busy)
+        self.assertEqual(self.stub.bodies[0]["priority"], 1)
+        chat.send()
+        chat.type("next question")  # an answer is streaming: no warm-up in "idle" mode
+        wait_for(lambda: chat.poll() or chat.reply is None)
+        self.assertEqual(chat.messages, [{"role": "user", "content": "Why is the sky"},
+                                         {"role": "assistant", "content": "Hello!"}])
+        self.assertEqual(len(self.stub.bodies), 2)
+        self.assertIn("tok/s", chat.status)
+
+    def test_warmup_off_and_cancel_on_send(self):
+        off = Chat(self.stub.url, "stub-model", warmup="off")
+        off.type("hello")
+        self.assertFalse(off.warmer.busy)
+        self.stub.delay = 3.0
+        chat = Chat(self.stub.url, "stub-model")
+        chat.type("a slow warm-up")
+        wait_for(lambda: self.stub.bodies)
+        chat.send()  # aborts the warm-up instead of queueing behind it
+        wait_for(lambda: not chat.warmer.busy, timeout=2.0)
+
+    def test_curses_screen_runs_in_a_real_terminal(self):
+        master, slave = pty.openpty()
+        fcntl.ioctl(slave, termios.TIOCSWINSZ, struct.pack("HHHH", 24, 80, 0, 0))
+        tui = subprocess.Popen([sys.executable, "-m", "prain.cli", "tui", "--url", self.stub.url],
+                               stdin=slave, stdout=slave, stderr=slave,
+                               env={**os.environ, "TERM": "xterm"}, start_new_session=True)
+        os.close(slave)
+        screen = bytearray()
+        threading.Thread(target=self._drain, args=(master, screen), daemon=True).start()
+        wait_for(lambda: b"stub-model" in screen)  # status line drawn
+        os.write(master, b"hi\r")
+        wait_for(lambda: any(body.get("stream") and "priority" not in body
+                             for body in self.stub.bodies))
+        wait_for(lambda: b"tok/s" in screen)
+        os.write(master, b"\x1b")
+        self.assertEqual(tui.wait(timeout=5), 0)
+
+    @staticmethod
+    def _drain(fd, screen):
+        while True:
+            try:
+                data = os.read(fd, 4096)
+            except OSError:
+                return
+            if not data:
+                return
+            screen.extend(data)