File size: 1,259 Bytes
3fd1a35 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 | #!/usr/bin/env python3
"""Prove a selected context is real, then verify reset after the long prompt."""
import argparse
import json
import urllib.request
p = argparse.ArgumentParser()
p.add_argument("--url", required=True)
p.add_argument("--digits", type=int, default=4200)
args = p.parse_args()
opener = urllib.request.build_opener(urllib.request.ProxyHandler({}))
def chat(text, count):
body = {"model": "mindnano-ling3-tiny", "messages": [{"role": "user", "content": text}],
"max_tokens": count, "temperature": 0}
req = urllib.request.Request(args.url + "/v1/chat/completions", json.dumps(body).encode(),
{"Content-Type": "application/json"})
with opener.open(req, timeout=300) as response:
return json.load(response)
a = chat("请只回答一个中文问候。", 16)
long = chat("这是测试编号，请只回答收到：" + "1" * args.digits, 1)
assert long["usage"]["prompt_tokens"] > 4096, long["usage"]
b = chat("请只回答一个中文问候。", 16)
assert a["choices"][0]["message"]["content"] == b["choices"][0]["message"]["content"], "long-context reset leaked state"
print(json.dumps({"passed": True, "long_request": long, "reset_reply": b}, ensure_ascii=False, indent=2))
|