fabolous005
a67238f66b
prain web/tui --persona FILE sends the file as the system message ahead of every warm-up and answer; the page shows its name. Three examples in examples/personas (task, how to answer, environment, how to use it). A persona costs 2.0-3.2 s once (241-467 tokens), then at most 0.3 s per conversation from the prefix cache. README rewritten around launching: quick start, the vLLM flags, a quantization table measured today at equal expert VRAM (bf16 0.90 / int8 2.19 / int4 5.19 tok/s; 2.3k-token prompt 69.5 / 32.3 / 20.5 s; 4 of 6 int8 and 0 of 6 int4 answers identical to bf16, one int4 answer wrong), how to switch formats, the VRAM budget, writing personas, the chats and the monitor.
82 lines
3.3 KiB
Python
82 lines
3.3 KiB
Python
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_a_persona_leads_every_request_but_stays_off_screen(self):
|
|
chat = Chat(self.stub.url, "stub-model", system="Be brief.")
|
|
chat.type("Why")
|
|
wait_for(lambda: self.stub.bodies and not chat.warmer.busy)
|
|
chat.send()
|
|
wait_for(lambda: chat.poll() or chat.reply is None)
|
|
system = {"role": "system", "content": "Be brief."}
|
|
self.assertEqual([b["messages"][0] for b in self.stub.bodies], [system, system])
|
|
self.assertEqual(chat.messages, [{"role": "user", "content": "Why"},
|
|
{"role": "assistant", "content": "Hello!"}])
|
|
|
|
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)
|