61 lines
2.8 KiB
Python
61 lines
2.8 KiB
Python
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
|