From 6e619b75d27e0ef69eb9056061a4f229e0d897cd Mon Sep 17 00:00:00 2001 From: Gahow Wang Date: Wed, 15 Jul 2026 19:53:00 +0800 Subject: [PATCH] Replay raw completion prompts without chat wrapping --- src/aituner/http_client.py | 12 ++++--- src/aituner/spec.py | 6 ++-- src/aituner/trace.py | 55 +++++++++++++++++++++----------- src/aituner/worker.py | 7 +++- tests/test_core_flow.py | 65 ++++++++++++++++++++++++++++++++++++++ 5 files changed, 118 insertions(+), 27 deletions(-) diff --git a/src/aituner/http_client.py b/src/aituner/http_client.py index 186cf0d..0adfa80 100644 --- a/src/aituner/http_client.py +++ b/src/aituner/http_client.py @@ -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: diff --git a/src/aituner/spec.py b/src/aituner/spec.py index a9c5720..a881e8e 100644 --- a/src/aituner/spec.py +++ b/src/aituner/spec.py @@ -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, diff --git a/src/aituner/trace.py b/src/aituner/trace.py index cabb3a5..9e39323 100644 --- a/src/aituner/trace.py +++ b/src/aituner/trace.py @@ -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"), diff --git a/src/aituner/worker.py b/src/aituner/worker.py index 0439537..d8f0ca6 100644 --- a/src/aituner/worker.py +++ b/src/aituner/worker.py @@ -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", "") diff --git a/tests/test_core_flow.py b/tests/test_core_flow.py index 465244a..1e9d0a6 100644 --- a/tests/test_core_flow.py +++ b/tests/test_core_flow.py @@ -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):