prain: terminal UI
This commit is contained in:
+115
@@ -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)
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user