diff --git a/prain/chat.py b/prain/chat.py new file mode 100644 index 0000000..f82a1d1 --- /dev/null +++ b/prain/chat.py @@ -0,0 +1,119 @@ +"""OpenAI-compatible client pieces shared by the TUI and the WebUI. Stdlib only.""" + +import http.client +import json +import socket +import threading +import urllib.parse +import urllib.request + +ABORTED = (OSError, ValueError, http.client.HTTPException) # what a cancelled request raises + + +def post(base_url, path, body, timeout=600): + request = urllib.request.Request(f"{base_url}{path}", json.dumps(body).encode(), + {"Content-Type": "application/json"}) + return urllib.request.urlopen(request, timeout=timeout) + + +def send_request(base_url, path, body, timeout=600): + """POST without waiting for the answer. Returns the connection: call getresponse() to + read it, or shutdown its socket from another thread to abort it (vLLM then cancels).""" + url = urllib.parse.urlsplit(base_url) + connection = http.client.HTTPConnection(url.hostname, url.port or 80, timeout=timeout) + connection.request("POST", path, json.dumps(body), {"Content-Type": "application/json"}) + return connection + + +def first_model(base_url): + with urllib.request.urlopen(f"{base_url}/v1/models", timeout=30) as response: + return json.load(response)["data"][0]["id"] + + +def chat_body(model, messages, **extra): + return {"model": model, "messages": messages, "stream": True, **extra} + + +def draft_body(model, messages): + """An answer computed from the experts already in VRAM: prain.drafts marks the request + by its vllm_xargs. Its own cache_salt keeps its KV blocks, computed with skipped experts, + out of the prefix-cache hits of full requests.""" + return chat_body(model, messages, vllm_xargs={"prain": "draft"}, cache_salt="prain-draft") + + +def warm_body(model, messages, draft, priority=0): + """Prefill the half-typed message and stop after one token. vLLM keeps its KV blocks + (prefix caching) and prain leaves the experts it routed to most in VRAM, so the real + request only prefills what was typed after this. Streamed, so aborting it is cheap. + priority > 0 queues it behind real requests (needs --scheduling-policy priority).""" + body = {"model": model, "messages": [*messages, {"role": "user", "content": draft.strip()}], + "max_tokens": 1, "temperature": 0, "stream": True, + "add_generation_prompt": False, "continue_final_message": True} + if priority: + body["priority"] = priority + return body + + +def stream_text(response): + """Yield the content pieces of a streamed chat completion (server-sent events).""" + for raw in response: + line = raw.decode().strip() + if not line.startswith("data: ") or line == "data: [DONE]": + continue + piece = json.loads(line[6:])["choices"][0].get("delta", {}).get("content") + if piece: + yield piece + + +class Warmer: + """Warm-up requests of one chat session. At most one is in flight; while it runs, newer + drafts replace each other and only the newest is sent next. cancel() drops the queued + draft and aborts the one in flight, so a sent message never waits behind a warm-up.""" + + def __init__(self, start): + self.start = start # start(job) -> connection from send_request() + self.lock = threading.Lock() + self.pending = None + self.busy = False + self.connection = None + self.sent = 0 + + def update(self, job): + with self.lock: + self.pending = job + if self.busy: + return + self.busy = True + threading.Thread(target=self._run, daemon=True).start() + + def cancel(self): + with self.lock: + self.pending = None + connection = self.connection + if connection is not None and connection.sock is not None: + try: + connection.sock.shutdown(socket.SHUT_RDWR) # wakes the blocked reader + except OSError: + pass + + def _run(self): + while True: + with self.lock: + job, self.pending = self.pending, None + if job is None: + self.busy = False + return + try: + connection = self.start(job) + with self.lock: + self.connection = connection + self.sent += 1 + with connection.getresponse() as response: + response.read() + except ABORTED: + pass # best effort: the real request reports server errors + finally: + with self.lock: + connection, self.connection = self.connection, None + if connection is not None: + connection.close() diff --git a/tests/test_chat.py b/tests/test_chat.py new file mode 100644 index 0000000..0d4dd7b --- /dev/null +++ b/tests/test_chat.py @@ -0,0 +1,60 @@ +import threading +import time +import unittest + +import prain.chat as chat +from tests.helpers import StubVllm, wait_for + + +class Chat(unittest.TestCase): + def setUp(self): + self.stub = StubVllm() + self.addCleanup(self.stub.close) + + def warmer(self): + return chat.Warmer(lambda draft: chat.send_request( + self.stub.url, "/v1/chat/completions", chat.warm_body("m", [], draft, priority=1))) + + def test_warm_body_continues_the_half_typed_message(self): + body = chat.warm_body("m", [{"role": "assistant", "content": "hi"}], " Explain mix ", 1) + self.assertEqual(body["messages"][-1], {"role": "user", "content": "Explain mix"}) + self.assertEqual( + (body["max_tokens"], body["add_generation_prompt"], body["continue_final_message"], + body["stream"], body["priority"]), (1, False, True, True, 1)) + self.assertNotIn("priority", chat.warm_body("m", [], "x")) + + def test_draft_body_marks_the_request_and_salts_its_cache(self): + body = chat.draft_body("m", [{"role": "user", "content": "hi"}]) + self.assertEqual(body, {"model": "m", "messages": [{"role": "user", "content": "hi"}], + "stream": True, "vllm_xargs": {"prain": "draft"}, + "cache_salt": "prain-draft"}) + + def test_first_model_and_streamed_text(self): + self.assertEqual(chat.first_model(self.stub.url), "stub-model") + with chat.post(self.stub.url, "/v1/chat/completions", + chat.chat_body("stub-model", [], max_tokens=5)) as response: + self.assertEqual(list(chat.stream_text(response)), ["Hel", "lo", "!"]) + self.assertEqual((self.stub.bodies[0]["stream"], self.stub.bodies[0]["max_tokens"]), + (True, 5)) + + def test_warmer_keeps_one_request_in_flight_and_sends_the_newest_draft(self): + self.stub.delay = 0.3 + warmer = self.warmer() + warmer.update("a") + wait_for(lambda: len(self.stub.bodies) == 1) + warmer.update("ab") # "a" is still in flight: these two queue up, + warmer.update("abc") # and only the newest is sent + wait_for(lambda: not warmer.busy) + self.assertEqual([b["messages"][-1]["content"] for b in self.stub.bodies], ["a", "abc"]) + + def test_cancel_drops_the_queue_and_aborts_the_request_in_flight(self): + self.stub.delay = 3.0 + warmer = self.warmer() + warmer.update("slow") + wait_for(lambda: self.stub.bodies) + warmer.update("queued") + start = time.monotonic() + warmer.cancel() + wait_for(lambda: not warmer.busy, timeout=2.0) + self.assertLess(time.monotonic() - start, 1.0) # did not wait out the 3 s answer + self.assertEqual(len(self.stub.bodies), 1) # "queued" was never sent