40 lines
2.2 KiB
Python
40 lines
2.2 KiB
Python
"""Check cached-prefix retrieval at several document positions, including changed suffixes."""
|
|
import json
|
|
import re
|
|
import urllib.request
|
|
from pathlib import Path
|
|
|
|
base = 'http://127.0.0.1:8000'
|
|
key = Path('/run/secrets/qwen_api_key').read_text().strip()
|
|
headers = {'Authorization': 'Bearer ' + key, 'Content-Type': 'application/json'}
|
|
def request(path, payload=None):
|
|
req = urllib.request.Request(base+path, headers=headers,
|
|
data=json.dumps(payload).encode() if payload is not None else None)
|
|
with urllib.request.urlopen(req, timeout=600) as response:
|
|
return response.read().decode()
|
|
def hits():
|
|
return sum(float(line.split()[-1]) for line in request('/metrics').splitlines()
|
|
if line.startswith('vllm:prefix_cache_hits_total'))
|
|
|
|
rows = [f'归档记录{i}：此行仅供背景阅读，无有效预算。' for i in range(700)]
|
|
for position, name, amount in [(20,'青松','17391'),(350,'白鹭','28647'),(670,'海棠','39583')]:
|
|
rows[position] = f'已核定：{name}项目预算为{amount}元。'
|
|
document = '\n'.join(rows)
|
|
before = hits()
|
|
for name, expected in [('青松','17391'),('白鹭','28647'),('海棠','39583'),('青松','17391')]:
|
|
# Only the question suffix changes. The document prefix is identical.
|
|
payload = {'model':'qwen3.8-flash-next', 'temperature':0, 'seed':42,
|
|
'reasoning_effort':'low','max_tokens':512,
|
|
'messages':[{'role':'user','content':'请根据以下档案回答，忽略无效归档行。\n<档案>\n'+document+'\n</档案>\n'+name+'项目的已核定预算是多少元？只输出数字。'}]}
|
|
result = json.loads(request('/v1/chat/completions',payload))
|
|
choice = result['choices'][0]
|
|
content = choice['message'].get('content') or ''
|
|
correct = content.strip() == expected and choice['finish_reason']=='stop'
|
|
print(json.dumps({'project':name,'expected':expected,'content':content,
|
|
'correct':correct,'usage':result.get('usage')},ensure_ascii=False),flush=True)
|
|
assert correct, 'Cached prefix answer mismatch'
|
|
delta = hits()-before
|
|
print(json.dumps({'prefix_cache_hits_delta':delta}),flush=True)
|
|
assert delta>0, 'No actual cache hits measured'
|
|
print('PREFIX_CHECK_PASS',flush=True)
|