refactor: recover per-token ids and logprobs from content[] logprobs shape

This client always requests boolean logprobs=True, so every response uses the
new content[] logprobs shape. Read the exact per-token ids and logprobs
together from content[] (token_id + sampling_logprob) so they are aligned by
construction.

Bug 1 (alignment / KLD corruption) — the real fix: previously ids were read
only from the top-level choices[].token_ids (populated by return_token_ids,
which this client never requested) and, when absent, silently re-derived via
tokenizer.encode(decode(text)). Retokenization drops the trailing
end-of-turn/EOS token, making completion_ids one shorter than the logprobs and
misaligning every per-token logprob (inference_kld ~60 vs ~0.028). Since
content[] already carries token_id per entry, reading ids there removes the
drift entirely — no return_token_ids, no re-encode fallback.

Bug 2 (precision) — more of a feature: prefer content[].sampling_logprob (the
exact value the sampler drew with) over content[].logprob (rounded).

- Drop legacy token_logprobs / top-level token_ids / raw_output id sources.
- Fail loud when content[] is absent or a content[] entry lacks token_id,
  instead of silently returning corrupted data.

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: Dranoxgithub

eval_protocol/integrations/fireworks_v1_completions_client.py

@@ -83,6 +83,19 @@ def _normalize_token_id_sequence(values: Any) -> List[int]:
     return [int(x) for x in list(values)]
 
 
+def _extract_entry_logprob(entry: Dict[str, Any]) -> float:
+    """Return the per-token logprob from a ``content[]`` logprobs entry.
+
+    Prefer ``sampling_logprob`` (the exact, full-precision value the sampler
+    actually drew with) over ``logprob`` (a rounded/verification value). The
+    sampler value is what RL training needs for accurate inference KLD.
+    """
+    value = entry.get("sampling_logprob")
+    if value is None:
+        value = entry.get("logprob", 0.0)
+    return float(value) if value is not None else 0.0
+
+
 def _coerce_message_content_to_text(content: Any) -> str:
     if content is None:
         return ""
@@ -358,12 +371,6 @@ async def create_completion_from_prompt_ids(
         }
         if not self.logprobs:
             request_payload.pop("logprobs", None)
-        # Always request the exact generated token ids so downstream RL training
-        # can align per-token logprobs to token ids position-by-position.
-        # Re-encoding decoded text drops the trailing end-of-turn/EOS token and
-        # misaligns logprobs, corrupting inference KLD. Keep overridable via
-        # request_params, but default the flag on.
-        request_payload.setdefault("return_token_ids", True)
 
         max_retries = 40
         base_delay = 10.0
@@ -399,52 +406,44 @@ async def create_completion_from_prompt_ids(
         finish_reason = str(choice.get("finish_reason") or "unknown")
 
         raw_output = choice.get("raw_output") if isinstance(choice.get("raw_output"), dict) else {}
-        completion_token_ids = _normalize_token_id_sequence(
-            choice.get("token_ids") or raw_output.get("completion_token_ids") or []
-        )
-        if not completion_token_ids:
-            raise RuntimeError(
-                "Fireworks /v1/completions returned no exact completion token IDs "
-                "(choices[].token_ids and raw_output.completion_token_ids both empty) "
-                "even though return_token_ids=True was requested. "
-                "Refusing to re-encode decoded text: retokenization drops the "
-                "end-of-turn token and misaligns per-token logprobs, corrupting "
-                f"inference KLD. choice keys={list(choice.keys())}"
-            )
         choice_prompt_token_ids = _normalize_token_id_sequence(
             choice.get("prompt_token_ids") or raw_output.get("prompt_token_ids") or normalized_prompt_token_ids
         )
 
-        completion_text = self.decode_token_ids(token_ids=completion_token_ids)
-        if not completion_text:
-            completion_text = str(choice.get("text") or "")
-
-        # -- Extract logprobs -----------------------------------------------
-        completion_logprobs: List[float] = []
+        # -- Extract per-token ids and logprobs together --------------------
+        # Both come from the same ``content[]`` array entry-by-entry, so they are
+        # inherently the same length and aligned. Reading ids from a different
+        # source (top-level token_ids) and re-encoding decoded text when it is
+        # absent drops the trailing end-of-turn token and misaligns per-token
+        # logprobs, corrupting inference KLD — so we never do that.
         choice_logprobs = choice.get("logprobs")
-        if isinstance(choice_logprobs, dict):
-            token_logprobs = choice_logprobs.get("token_logprobs") or []
-            if token_logprobs:
-                completion_logprobs = [float(lp) if lp is not None else 0.0 for lp in token_logprobs]
-            else:
-                content_logprobs = choice_logprobs.get("content") or []
-                completion_logprobs = [
-                    float(entry.get("logprob", 0.0)) if isinstance(entry, dict) else 0.0
-                    for entry in content_logprobs
-                ]
-        elif isinstance(choice_logprobs, list):
-            completion_logprobs = [float(lp) if lp is not None else 0.0 for lp in choice_logprobs]
-
-        # Catch any residual token-id / logprob drift at the boundary rather than
-        # as a downstream KLD anomaly.
-        if completion_logprobs and len(completion_token_ids) != len(completion_logprobs):
+        content_entries = choice_logprobs.get("content") if isinstance(choice_logprobs, dict) else None
+        if not content_entries:
             raise RuntimeError(
-                "Fireworks /v1/completions returned mismatched completion token "
-                f"ids ({len(completion_token_ids)}) and logprobs "
-                f"({len(completion_logprobs)}). Per-token logprobs cannot be "
-                "aligned to token ids; refusing to return corrupted data."
+                "Fireworks /v1/completions returned no content[] logprobs entries. "
+                "This client requires the boolean logprobs=True (content) shape to "
+                "recover exact per-token ids and sampling logprobs. Refusing to "
+                "re-encode decoded text: retokenization drops the end-of-turn token "
+                "and misaligns per-token logprobs, corrupting inference KLD. "
+                f"choice keys={list(choice.keys())}"
             )
 
+        completion_token_ids: List[int] = []
+        completion_logprobs: List[float] = []
+        for index, entry in enumerate(content_entries):
+            if not isinstance(entry, dict) or entry.get("token_id") is None:
+                raise RuntimeError(
+                    "Fireworks /v1/completions content[] entry is missing token_id "
+                    f"at index {index}; cannot align per-token logprobs to token ids "
+                    "without re-encoding. Refusing to return corrupted data."
+                )
+            completion_token_ids.append(int(entry["token_id"]))
+            completion_logprobs.append(_extract_entry_logprob(entry))
+
+        completion_text = self.decode_token_ids(token_ids=completion_token_ids)
+        if not completion_text:
+            completion_text = str(choice.get("text") or "")
+
         # -- Build message via parser or raw --------------------------------
         if self.tool_call_parser is not None:
             parsed_output = self.tool_call_parser(completion_text, completion_token_ids, active_tools)

tests/test_fireworks_v1_completions_client.py

@@ -152,95 +152,86 @@ async def close(self):
     return captured
 
 
-def test_request_payload_sets_return_token_ids(monkeypatch):
+def test_reads_ids_and_sampling_logprobs_from_content(monkeypatch):
     client = FireworksV1CompletionsClient(
         model_id="test-model",
         tokenizer_name_or_path="Qwen/Qwen3-0.6B",
     )
-    monkeypatch.setattr(client, "decode_token_ids", lambda token_ids: "hello")
+    monkeypatch.setattr(client, "decode_token_ids", lambda token_ids: "text")
+
+    def _fail_tokenizer():
+        raise AssertionError("tokenizer must not be used to re-encode completion text")
+
+    monkeypatch.setattr(client, "_get_tokenizer", lambda: _fail_tokenizer())
     captured = _install_fake_completion(
         client,
         monkeypatch,
         {
             "choices": [
                 {
-                    "token_ids": [5, 6, 7],
                     "finish_reason": "stop",
-                    "logprobs": {"token_logprobs": [-0.1, -0.2, -0.3]},
+                    "logprobs": {
+                        "content": [
+                            {"token_id": 271, "sampling_logprob": -0.05483185, "logprob": -0.0548313},
+                            {"token_id": 248068, "sampling_logprob": -0.0014, "logprob": -0.0014},
+                            {"token_id": 26108, "logprob": -1.0},
+                        ]
+                    },
                 }
             ],
         },
     )
     result = asyncio.run(client.create_completion_from_prompt_ids(prompt_token_ids=[1, 2]))
-    assert captured["return_token_ids"] is True
-    assert "raw_output" not in captured
-    assert result["completion_ids"] == [5, 6, 7]
+    assert "return_token_ids" not in captured
+    assert result["completion_ids"] == [271, 248068, 26108]
+    assert result["completion_logprobs"] == [-0.05483185, -0.0014, -1.0]
     assert len(result["completion_ids"]) == len(result["completion_logprobs"])
     asyncio.run(client.close())
 
 
-def test_request_params_can_override_flags(monkeypatch):
+def test_raises_when_no_content_logprobs(monkeypatch):
     client = FireworksV1CompletionsClient(
         model_id="test-model",
         tokenizer_name_or_path="Qwen/Qwen3-0.6B",
-        request_params={"return_token_ids": False},
     )
-    monkeypatch.setattr(client, "decode_token_ids", lambda token_ids: "hi")
-    captured = _install_fake_completion(
+    _install_fake_completion(
         client,
         monkeypatch,
-        {"choices": [{"token_ids": [9], "finish_reason": "stop"}]},
+        {"choices": [{"text": "hello world", "finish_reason": "stop"}]},
     )
-    asyncio.run(client.create_completion_from_prompt_ids(prompt_token_ids=[1]))
-    assert captured["return_token_ids"] is False
+    with pytest.raises(RuntimeError, match="no content\\[\\] logprobs entries"):
+        asyncio.run(client.create_completion_from_prompt_ids(prompt_token_ids=[1, 2]))
     asyncio.run(client.close())
 
 
-def test_uses_exact_token_ids_without_reencode(monkeypatch):
+def test_ignores_legacy_token_logprobs_shape(monkeypatch):
+    """The legacy token_logprobs shape has no content[]; the client must not use it."""
     client = FireworksV1CompletionsClient(
         model_id="test-model",
         tokenizer_name_or_path="Qwen/Qwen3-0.6B",
     )
-    monkeypatch.setattr(client, "decode_token_ids", lambda token_ids: "text")
-
-    def _fail_encode():
-        raise AssertionError("tokenizer must not be used to re-encode completion text")
-
-    monkeypatch.setattr(client, "_get_tokenizer", lambda: _fail_encode())
     _install_fake_completion(
         client,
         monkeypatch,
         {
             "choices": [
                 {
-                    "token_ids": [10, 20, 30, 40],
                     "finish_reason": "stop",
-                    "logprobs": {"token_logprobs": [-0.1, -0.2, -0.3, -0.4]},
+                    "token_ids": [1, 2, 3],
+                    "logprobs": {
+                        "token_ids": [1, 2, 3],
+                        "token_logprobs": [-0.1, -0.2, -0.3],
+                    },
                 }
             ],
         },
     )
-    result = asyncio.run(client.create_completion_from_prompt_ids(prompt_token_ids=[1]))
-    assert result["completion_ids"] == [10, 20, 30, 40]
-    asyncio.run(client.close())
-
-
-def test_raises_when_no_exact_token_ids(monkeypatch):
-    client = FireworksV1CompletionsClient(
-        model_id="test-model",
-        tokenizer_name_or_path="Qwen/Qwen3-0.6B",
-    )
-    _install_fake_completion(
-        client,
-        monkeypatch,
-        {"choices": [{"text": "hello world", "finish_reason": "stop"}]},
-    )
-    with pytest.raises(RuntimeError, match="no exact completion token IDs"):
-        asyncio.run(client.create_completion_from_prompt_ids(prompt_token_ids=[1, 2]))
+    with pytest.raises(RuntimeError, match="no content\\[\\] logprobs entries"):
+        asyncio.run(client.create_completion_from_prompt_ids(prompt_token_ids=[1]))
     asyncio.run(client.close())
 
 
-def test_raises_on_id_logprob_length_mismatch(monkeypatch):
+def test_raises_when_content_entry_missing_token_id(monkeypatch):
     client = FireworksV1CompletionsClient(
         model_id="test-model",
         tokenizer_name_or_path="Qwen/Qwen3-0.6B",
@@ -252,14 +243,18 @@ def test_raises_on_id_logprob_length_mismatch(monkeypatch):
         {
             "choices": [
                 {
-                    "token_ids": [1, 2, 3],
                     "finish_reason": "stop",
-                    "logprobs": {"token_logprobs": [-0.1, -0.2, -0.3, -0.4]},
+                    "logprobs": {
+                        "content": [
+                            {"token_id": 1, "sampling_logprob": -0.1},
+                            {"sampling_logprob": -0.2},
+                        ]
+                    },
                 }
             ],
         },
     )
-    with pytest.raises(RuntimeError, match="mismatched completion token"):
+    with pytest.raises(RuntimeError, match="missing token_id"):
         asyncio.run(client.create_completion_from_prompt_ids(prompt_token_ids=[1]))
     asyncio.run(client.close())