55 lines
2.0 KiB
Python
55 lines
2.0 KiB
Python
"""A second answer that arrives while a first one streams, through `prain web`: when does
|
|
its first token come, and how much of its draft survives the refine?
|
|
|
|
Run it against a proxy with drafts and one without ([draft] full_requests = max_requests)
|
|
and compare. usage: python3 scripts/draft_effect.py http://127.0.0.1:8080
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import sys
|
|
import threading
|
|
import time
|
|
|
|
from prain.chat import post
|
|
|
|
LONG = [{"role": "user", "content": "Write a detailed 400-word history of the printing press."}]
|
|
SHORT = [{"role": "user", "content": "Name three uses of copper, one line each."}]
|
|
|
|
|
|
def answer(url, session, messages, out):
|
|
start, mode = time.monotonic(), "full"
|
|
out.update(first=None, draft="", final="")
|
|
with post(url, "/chat", {"session": session, "messages": messages}) as response:
|
|
for raw in response:
|
|
line = raw.decode().strip()
|
|
if not line.startswith("data: ") or line == "data: [DONE]":
|
|
continue
|
|
data = json.loads(line[6:])
|
|
if "prain" in data:
|
|
mode = data["prain"]
|
|
continue
|
|
piece = data["choices"][0]["delta"].get("content")
|
|
if piece:
|
|
out["first"] = out["first"] or time.monotonic() - start
|
|
out["draft" if mode == "draft" else "final"] += piece
|
|
out["total"] = time.monotonic() - start
|
|
|
|
|
|
def main():
|
|
url = sys.argv[1]
|
|
first, second = {}, {}
|
|
thread = threading.Thread(target=answer, args=(url, "first", LONG, first))
|
|
thread.start()
|
|
time.sleep(5) # the first answer is decoding now
|
|
answer(url, "second", SHORT, second)
|
|
thread.join()
|
|
kept = len(os.path.commonprefix([second["draft"], second["final"]]))
|
|
print(f"second answer: first token {second['first']:.1f} s, done after {second['total']:.1f} s; "
|
|
+ (f"draft {len(second['draft'])} chars, {kept} kept by the refine"
|
|
if second["draft"] else "no draft"))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|