From 0531aaed44e4685ca41a6212d44deac4ca9081db Mon Sep 17 00:00:00 2001 From: fabolous005 Date: Tue, 27 Oct 2026 13:52:16 +0100 Subject: [PATCH] prain: terminal UI --- prain/tui.py | 115 ++++++++++++++++++++++++++++++++++++++++++++++ tests/test_tui.py | 70 ++++++++++++++++++++++++++++ 2 files changed, 185 insertions(+) create mode 100644 prain/tui.py create mode 100644 tests/test_tui.py diff --git a/prain/tui.py b/prain/tui.py new file mode 100644 index 0000000..1d4d002 --- /dev/null +++ b/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) diff --git a/tests/test_tui.py b/tests/test_tui.py new file mode 100644 index 0000000..51fe063 --- /dev/null +++ b/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)