diff --git a/prain/web.py b/prain/web.py
new file mode 100644
index 0000000..4e3040e
--- /dev/null
+++ b/prain/web.py
@@ -0,0 +1,274 @@
+"""prain web: a one-page browser chat that warms experts up while you type.
+
+A small stdlib proxy in front of vLLM. GET / serves the page. POST /warm takes
+{"session", "messages", "draft"} and returns at once. POST /chat takes {"session",
+"messages"} and streams server-sent events: vLLM's chunks, plus {"prain": "draft" |
+"refine" | "error"} events of its own. Binds to 127.0.0.1; use an SSH tunnel from other
+machines.
+
+Several browsers at once ([serve] and [draft] tables of the placement file):
+- At most draft.full_requests answers run with all their experts. An answer that arrives
+ while they are taken runs as a draft right away, on the experts already in VRAM (see
+ prain.drafts), instead of waiting or evicting the other answers' experts.
+- With draft.refine = "regenerate", the oldest waiting draft takes the next free full slot:
+ its draft request is cancelled and the answer is generated again with all experts. The
+ page shows the draft in grey and the full answer over it.
+- Every session has its own Warmer. warmup="idle": no warm-ups while any answer streams;
+ "always": warm-ups while vLLM has a free request slot (serve.max_requests). Sending cancels
+ that session's warm-up. A browser that goes away closes the upstream request, and vLLM
+ cancels it.
+"""
+
+import json
+import threading
+import urllib.error
+from collections import OrderedDict, deque
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+
+from prain.chat import Warmer, chat_body, draft_body, first_model, post, send_request, warm_body
+from prain.config import Draft, Serve
+
+MAX_SESSIONS = 64
+
+PAGE = """
+
+
prain
+
+
+
+
+
+"""
+
+
+def make_server(base_url, port, host="127.0.0.1", warmup="idle", priority=1, serve=Serve(),
+ draft=Draft()):
+ model = first_model(base_url)
+ lock = threading.Condition()
+ sessions = OrderedDict() # session -> Warmer, least recently used first
+ streams = [0] # answers streaming right now, all sessions, full or draft
+ full = [0] # answers holding one of the draft.full_requests full slots
+ waiting = deque() # drafts waiting for a full slot to refine, oldest first
+ refine = draft.refine == "regenerate"
+
+ def warmer(session):
+ with lock:
+ if session not in sessions:
+ sessions[session] = Warmer(lambda job: send_request(
+ base_url, "/v1/chat/completions", warm_body(model, *job, priority)))
+ if len(sessions) > MAX_SESSIONS:
+ sessions.popitem(last=False)[1].cancel()
+ sessions.move_to_end(session)
+ return sessions[session]
+
+ def my_turn(ticket): # call with the lock held
+ return full[0] < draft.full_requests and bool(waiting) and waiting[0] is ticket
+
+ def turn_now(ticket):
+ with lock:
+ return my_turn(ticket)
+
+ class Handler(BaseHTTPRequestHandler):
+ def do_GET(self):
+ self.reply(200, "text/html; charset=utf-8", PAGE.encode())
+
+ def do_POST(self):
+ body = json.loads(self.rfile.read(int(self.headers["Content-Length"])))
+ session = str(body.get("session", ""))
+ if self.path == "/warm":
+ with lock:
+ busy = streams[0]
+ if body["draft"].strip() and ((warmup == "idle" and not busy) or
+ (warmup == "always" and busy < serve.max_requests)):
+ warmer(session).update((body["messages"], body["draft"]))
+ return self.reply(204, "text/plain", b"")
+ if self.path != "/chat":
+ return self.reply(404, "text/plain", b"not found")
+ warmer(session).cancel()
+ messages, ticket = body["messages"], object()
+ with lock:
+ streams[0] += 1
+ holds_full = full[0] < draft.full_requests and not waiting
+ if holds_full:
+ full[0] += 1
+ elif refine:
+ waiting.append(ticket)
+ try:
+ if holds_full:
+ if not self.relay(chat_body(model, messages), opens=True):
+ return
+ else:
+ if not self.relay(draft_body(model, messages), opens=True, event="draft",
+ stop=(lambda: turn_now(ticket)) if refine else None):
+ return
+ if refine:
+ with lock:
+ lock.wait_for(lambda: my_turn(ticket))
+ waiting.popleft()
+ full[0] += 1
+ holds_full = True
+ self.relay(chat_body(model, messages), opens=False, event="refine")
+ self.send_event("[DONE]")
+ except (BrokenPipeError, ConnectionResetError):
+ pass # the browser left; closing upstream makes vLLM cancel the request
+ finally:
+ with lock:
+ streams[0] -= 1
+ if holds_full:
+ full[0] -= 1
+ if ticket in waiting:
+ waiting.remove(ticket)
+ lock.notify_all()
+
+ def relay(self, body, opens, event=None, stop=None):
+ """Stream one vLLM answer to the browser, without its [DONE]. opens: this starts
+ the response, so a vLLM error becomes the HTTP status (returns False). stop():
+ checked after every chunk; True closes the upstream request early."""
+ try:
+ upstream = post(base_url, "/v1/chat/completions", body)
+ except urllib.error.HTTPError as e: # vLLM rejected it: pass its answer on
+ with e:
+ data = e.read()
+ if opens:
+ self.reply(e.code, "application/json", data)
+ return False
+ self.send_event({"prain": "error", "status": e.code})
+ return True
+ with upstream:
+ if opens:
+ self.send_response(200)
+ self.send_header("Content-Type", "text/event-stream")
+ self.end_headers()
+ if event:
+ self.send_event({"prain": event})
+ for line in upstream:
+ if line.strip() in (b"", b"data: [DONE]"):
+ continue
+ self.wfile.write(line.rstrip(b"\n") + b"\n\n")
+ self.wfile.flush()
+ if stop is not None and stop():
+ break
+ return True
+
+ def send_event(self, data):
+ text = data if isinstance(data, str) else json.dumps(data)
+ self.wfile.write(f"data: {text}\n\n".encode())
+ self.wfile.flush()
+
+ def reply(self, code, kind, data):
+ self.send_response(code)
+ self.send_header("Content-Type", kind)
+ self.send_header("Content-Length", str(len(data)))
+ self.end_headers()
+ self.wfile.write(data)
+
+ def log_message(self, *args):
+ pass
+
+ server = ThreadingHTTPServer((host, port), Handler)
+ server.sessions, server.streams, server.full = sessions, streams, full # for tests
+ return server
+
+
+def serve(base_url, port, warmup="idle", priority=1, serve=Serve(), draft=Draft()):
+ server = make_server(base_url, port, warmup=warmup, priority=priority, serve=serve, draft=draft)
+ print(f"prain web: http://127.0.0.1:{server.server_port} (vLLM at {base_url}; "
+ f"{draft.full_requests} full of {serve.max_requests} parallel answers)")
+ server.serve_forever()
diff --git a/tests/test_web.py b/tests/test_web.py
new file mode 100644
index 0000000..24aa835
--- /dev/null
+++ b/tests/test_web.py
@@ -0,0 +1,138 @@
+import json
+import threading
+import unittest
+import urllib.error
+import urllib.request
+
+import prain.chat as chat
+from prain.config import Draft, Serve
+from prain.web import make_server
+from tests.helpers import StubVllm, wait_for
+
+
+def events(response):
+ """The proxy's stream as a list: content pieces, and ""-style markers."""
+ out = []
+ for raw in response:
+ line = raw.decode().strip()
+ if line == "data: [DONE]":
+ out.append("")
+ elif line.startswith("data: "):
+ data = json.loads(line[6:])
+ out.append(f"<{data['prain']}>" if "prain" in data
+ else data["choices"][0]["delta"]["content"])
+ return out
+
+
+class Web(unittest.TestCase):
+ def setUp(self):
+ self.stub = StubVllm()
+ self.addCleanup(self.stub.close)
+ self.start()
+
+ def start(self, **options):
+ self.server = make_server(self.stub.url, 0, **options)
+ threading.Thread(target=self.server.serve_forever, daemon=True).start()
+ self.addCleanup(self.server.server_close)
+ self.addCleanup(self.server.shutdown)
+ self.url = f"http://127.0.0.1:{self.server.server_port}"
+
+ def warm(self, session, draft):
+ chat.post(self.url, "/warm", {"session": session, "messages": [], "draft": draft}).read()
+
+ def ask(self, session, text):
+ return chat.post(self.url, "/chat", {"session": session,
+ "messages": [{"role": "user", "content": text}]})
+
+ def test_page_warm_up_and_chat(self):
+ with urllib.request.urlopen(self.url) as response:
+ self.assertIn(b'