fix: request exact token ids and fail loud in FireworksV1CompletionsClient

eval_protocol/integrations/fireworks_v1_completions_client.py

@@ -358,6 +358,12 @@ 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
@@ -396,16 +402,22 @@ async def create_completion_from_prompt_ids(
         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 "")
-        if not completion_token_ids and completion_text:
-            tokenizer = self._get_tokenizer()
-            completion_token_ids = list(tokenizer.encode(completion_text, add_special_tokens=False))
 
         # -- Extract logprobs -----------------------------------------------
         completion_logprobs: List[float] = []
@@ -423,6 +435,16 @@ async def create_completion_from_prompt_ids(
         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):
+            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."
+            )
+
         # -- 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

@@ -124,6 +124,146 @@ def encode(self, text, add_special_tokens=False):
     asyncio.run(client.close())
 
 
+class _FakeResponse:
+    def __init__(self, payload: Dict[str, Any]):
+        self._payload = payload
+
+    def model_dump(self) -> Dict[str, Any]:
+        return self._payload
+
+
+def _install_fake_completion(client, monkeypatch, payload):
+    captured: Dict[str, Any] = {}
+
+    async def fake_create(**kwargs):
+        captured.update(kwargs)
+        return _FakeResponse(payload)
+
+    class _FakeCompletions:
+        create = staticmethod(fake_create)
+
+    class _FakeClient:
+        completions = _FakeCompletions()
+
+        async def close(self):
+            return None
+
+    monkeypatch.setattr(client, "_client", _FakeClient())
+    return captured
+
+
+def test_request_payload_sets_return_token_ids(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")
+    captured = _install_fake_completion(
+        client,
+        monkeypatch,
+        {
+            "choices": [
+                {
+                    "token_ids": [5, 6, 7],
+                    "finish_reason": "stop",
+                    "logprobs": {"token_logprobs": [-0.1, -0.2, -0.3]},
+                }
+            ],
+        },
+    )
+    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 len(result["completion_ids"]) == len(result["completion_logprobs"])
+    asyncio.run(client.close())
+
+
+def test_request_params_can_override_flags(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(
+        client,
+        monkeypatch,
+        {"choices": [{"token_ids": [9], "finish_reason": "stop"}]},
+    )
+    asyncio.run(client.create_completion_from_prompt_ids(prompt_token_ids=[1]))
+    assert captured["return_token_ids"] is False
+    asyncio.run(client.close())
+
+
+def test_uses_exact_token_ids_without_reencode(monkeypatch):
+    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]},
+                }
+            ],
+        },
+    )
+    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]))
+    asyncio.run(client.close())
+
+
+def test_raises_on_id_logprob_length_mismatch(monkeypatch):
+    client = FireworksV1CompletionsClient(
+        model_id="test-model",
+        tokenizer_name_or_path="Qwen/Qwen3-0.6B",
+    )
+    monkeypatch.setattr(client, "decode_token_ids", lambda token_ids: "text")
+    _install_fake_completion(
+        client,
+        monkeypatch,
+        {
+            "choices": [
+                {
+                    "token_ids": [1, 2, 3],
+                    "finish_reason": "stop",
+                    "logprobs": {"token_logprobs": [-0.1, -0.2, -0.3, -0.4]},
+                }
+            ],
+        },
+    )
+    with pytest.raises(RuntimeError, match="mismatched completion token"):
+        asyncio.run(client.create_completion_from_prompt_ids(prompt_token_ids=[1]))
+    asyncio.run(client.close())
+
+
 def test_thinking_kwargs_respects_enable_thinking():
     client_none = FireworksV1CompletionsClient(
         model_id="test", tokenizer_name_or_path="Qwen/Qwen3-0.6B",