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", {})
-                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:

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,

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

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