gahow/aituner
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Replay raw completion prompts without chat wrapping

Browse Source
This commit is contained in:
Gahow Wang
2026-07-15 19:53:00 +08:00
parent 584af7b253
commit 6e619b75d2
5 changed files with 118 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
src/aituner/http_client.py
View File

@@ -254,10 +254,11 @@ def stream_chat_completion(
base_url: str,
body: dict[str, Any],
timeout_s: float,
api_path: str = "/v1/chat/completions",
) -> StreamMetrics:
data = json.dumps(body).encode("utf-8")
request = urllib.request.Request(
url=_openai_url(base_url, "/v1/chat/completions"),
url=_openai_url(base_url, api_path),
headers=_auth_headers(None),
data=data,
method="POST",
@@ -285,10 +286,11 @@ def stream_chat_completion(
choices = payload.get("choices")
if not isinstance(choices, list) or not choices:
continue
delta = choices[0].get("delta", {})
if not isinstance(delta, dict):
continue
content = delta.get("content")
choice = choices[0]
delta = choice.get("delta", {})
content = delta.get("content") if isinstance(delta, dict) else None
if not isinstance(content, str):
content = choice.get("text")
if isinstance(content, str) and content:
now = time.monotonic()
if first_token_at is None:

6
src/aituner/spec.py
View File

@@ -411,8 +411,10 @@ class TraceSpec:
synthetic_prompt_cap = data.get("synthetic_prompt_cap_tokens")
completion_tokens_override = data.get("completion_tokens_override")
request_mode = str(data.get("request_mode") or "chat").strip().lower()
if request_mode not in {"chat", "decode_only"}:
raise SpecError("trace.request_mode must be one of: chat, decode_only.")
if request_mode not in {"chat", "decode_only", "raw_completion"}:
raise SpecError(
"trace.request_mode must be one of: chat, decode_only, raw_completion."
)
if completion_tokens_override is not None:
completion_tokens_override = _require_int(
completion_tokens_override,

55
src/aituner/trace.py
View File

@@ -40,6 +40,7 @@ class TraceRequest:
body: dict[str, Any]
prompt_tokens_hint: int | None
completion_tokens_hint: int | None
api_path: str = "/v1/chat/completions"
metadata: dict[str, Any] = field(default_factory=dict)
@@ -186,26 +187,41 @@ def load_trace_requests(study: StudySpec, *, study_spec_path: Path) -> tuple[Win
prompt_tokens_hint = _coerce_prompt_tokens(row)
if not _matches_input_length_filter(study, prompt_tokens_hint=prompt_tokens_hint):
continue
try:
messages = _coerce_messages(row)
except TraceError:
capped_prompt_tokens = prompt_tokens_hint or 0
if study.trace.synthetic_prompt_cap_tokens is not None:
capped_prompt_tokens = min(
capped_prompt_tokens, study.trace.synthetic_prompt_cap_tokens
api_path = "/v1/chat/completions"
if study.trace.request_mode == "raw_completion":
prompt = row.get("prompt")
if not isinstance(prompt, str) or not prompt:
raise TraceError(
f"trace row {idx} is missing prompt required by raw_completion"
)
messages = [
{
"role": "user",
"content": _synthetic_prompt_from_tokens(capped_prompt_tokens),
}
]
body: dict[str, Any] = {
"model": study.model.served_model_name,
"messages": messages,
"stream": True,
"stream_options": {"include_usage": True},
}
body: dict[str, Any] = {
"model": study.model.served_model_name,
"prompt": prompt,
"stream": True,
"stream_options": {"include_usage": True},
}
api_path = "/v1/completions"
else:
try:
messages = _coerce_messages(row)
except TraceError:
capped_prompt_tokens = prompt_tokens_hint or 0
if study.trace.synthetic_prompt_cap_tokens is not None:
capped_prompt_tokens = min(
capped_prompt_tokens, study.trace.synthetic_prompt_cap_tokens
)
messages = [
{
"role": "user",
"content": _synthetic_prompt_from_tokens(capped_prompt_tokens),
}
]
body = {
"model": study.model.served_model_name,
"messages": messages,
"stream": True,
"stream_options": {"include_usage": True},
}
completion_tokens = (
study.trace.completion_tokens_override
if study.trace.completion_tokens_override is not None
@@ -225,6 +241,7 @@ def load_trace_requests(study: StudySpec, *, study_spec_path: Path) -> tuple[Win
body=body,
prompt_tokens_hint=prompt_tokens_hint,
completion_tokens_hint=completion_tokens,
api_path=api_path,
metadata={
"hash_ids": row.get("hash_ids") if isinstance(row.get("hash_ids"), list) else None,
"turn": row.get("turn"),

7
src/aituner/worker.py
View File

@@ -107,7 +107,12 @@ def _run_one_request(
timeout_s: float,
) -> RequestOutcome:
try:
metrics = stream_chat_completion(base_url=base_url, body=request.body, timeout_s=timeout_s)
metrics = stream_chat_completion(
base_url=base_url,
body=request.body,
timeout_s=timeout_s,
api_path=request.api_path,
)
expected_completion_tokens = request.completion_tokens_hint
actual_completion_tokens = metrics.completion_tokens
completion_tokens_source = getattr(metrics, "completion_tokens_source", "")

65
tests/test_core_flow.py
View File

@@ -5586,6 +5586,42 @@ class CoreFlowTests(unittest.TestCase):
self.assertEqual(requests[2].body["min_tokens"], 1)
self.assertEqual(requests[2].body["max_tokens"], 1)
def test_raw_completion_mode_preserves_trace_prompt_and_endpoint(self) -> None:
with tempfile.TemporaryDirectory() as tmp:
tmp_path = Path(tmp)
study_path = _write_study_assets(
tmp_path,
trace_overrides={"request_mode": "raw_completion"},
)
trace_path = tmp_path / "trace_windows" / "traces" / "chat_w1.jsonl"
rows = [json.loads(line) for line in trace_path.read_text().splitlines()]
for idx, row in enumerate(rows):
row["prompt"] = f"<|im_start|>user\nraw prompt {idx}<|im_end|>"
trace_path.write_text(
"".join(json.dumps(row) + "\n" for row in rows),
encoding="utf-8",
)
study = load_study_spec(study_path)
_, requests = load_trace_requests(study, study_spec_path=study_path)
self.assertEqual(study.trace.request_mode, "raw_completion")
self.assertEqual(requests[0].api_path, "/v1/completions")
self.assertEqual(requests[0].body["prompt"], rows[0]["prompt"])
self.assertNotIn("messages", requests[0].body)
def test_raw_completion_mode_requires_prompt(self) -> None:
with tempfile.TemporaryDirectory() as tmp:
study_path = _write_study_assets(
Path(tmp),
trace_overrides={"request_mode": "raw_completion"},
)
study = load_study_spec(study_path)
with self.assertRaisesRegex(
ValueError, "missing prompt required by raw_completion"
):
load_trace_requests(study, study_spec_path=study_path)
def test_run_one_request_fails_fixed_length_completion_mismatch(self) -> None:
request = TraceRequest(
row_id="r1",
@@ -8815,6 +8851,35 @@ class CoreFlowTests(unittest.TestCase):
self.assertIsNone(metrics.completion_tokens)
self.assertEqual(metrics.completion_tokens_source, "none")
def test_stream_chat_completion_reads_raw_completion_text(self) -> None:
class FakeResponse:
def __enter__(self):
return self
def __exit__(self, exc_type, exc, traceback):
return False
def __iter__(self):
return iter(
[
b'data: {"choices": [{"text": "x"}]}\n',
b'data: {"choices": [], "usage": {"completion_tokens": 1}}\n',
b"data: [DONE]\n",
]
)
with mock.patch("aituner.http_client._urlopen", return_value=FakeResponse()):
metrics = stream_chat_completion(
base_url="http://127.0.0.1:8000",
body={"model": "m", "prompt": "raw"},
timeout_s=1.0,
api_path="/v1/completions",
)
self.assertIsNotNone(metrics.ttft_ms)
self.assertEqual(metrics.completion_tokens, 1)
self.assertEqual(metrics.streamed_chunk_count, 1)
def test_stream_chat_completion_marks_same_instant_multitoken_tpot_unmeasurable(self) -> None:
class FakeResponse:
def __enter__(self):