Replay raw completion prompts without chat wrapping

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", {})
+                choice = choices[0]
-                if not isinstance(delta, dict):
+                delta = choice.get("delta", {})
-                    continue
+                content = delta.get("content") if isinstance(delta, dict) else None
-                content = delta.get("content")
+                if not isinstance(content, str):
+                    content = choice.get("text")
                 if isinstance(content, str) and content:
                     now = time.monotonic()
                     if first_token_at is None:

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"}:
+        if request_mode not in {"chat", "decode_only", "raw_completion"}:
-            raise SpecError("trace.request_mode must be one of: chat, decode_only.")
+            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,

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:
+            api_path = "/v1/chat/completions"
-                messages = _coerce_messages(row)
+            if study.trace.request_mode == "raw_completion":
-            except TraceError:
+                prompt = row.get("prompt")
-                capped_prompt_tokens = prompt_tokens_hint or 0
+                if not isinstance(prompt, str) or not prompt:
-                if study.trace.synthetic_prompt_cap_tokens is not None:
+                    raise TraceError(
-                    capped_prompt_tokens = min(
+                        f"trace row {idx} is missing prompt required by raw_completion"
-                        capped_prompt_tokens, study.trace.synthetic_prompt_cap_tokens
                     )
-                messages = [
+                body: dict[str, Any] = {
-                    {
+                    "model": study.model.served_model_name,
-                        "role": "user",
+                    "prompt": prompt,
-                        "content": _synthetic_prompt_from_tokens(capped_prompt_tokens),
+                    "stream": True,
-                    }
+                    "stream_options": {"include_usage": True},
-                ]
+                }
-            body: dict[str, Any] = {
+                api_path = "/v1/completions"
-                "model": study.model.served_model_name,
+            else:
-                "messages": messages,
+                try:
-                "stream": True,
+                    messages = _coerce_messages(row)
-                "stream_options": {"include_usage": True},
+                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"),

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", "")

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):