From 64356d0bec9cc06f286b175ac715e5ee6a17b9a5 Mon Sep 17 00:00:00 2001 From: Anni Chen Date: Sun, 5 Jul 2026 22:13:19 -0700 Subject: [PATCH 1/4] fix: request exact token ids and fail loud in FireworksV1CompletionsClient The /v1/completions sampling path corrupted RL training via retokenization drift. The request never set return_token_ids, so choices[].token_ids came back empty and the client silently re-encoded decoded text. Retokenization drops the trailing end-of-turn/EOS token, making completion_ids one shorter than token_logprobs and misaligning every per-token logprob (inference KLD spiked to ~60 vs ~0.028 with exact ids). - Default return_token_ids=True in the request payload (overridable via request_params). - Remove the tokenizer.encode() re-encode fallback; raise when exact ids are absent instead of silently returning corrupted data. - Assert len(completion_token_ids) == len(completion_logprobs) at the boundary to catch any residual drift. Co-authored-by: Cursor --- .../fireworks_v1_completions_client.py | 28 +++- tests/test_fireworks_v1_completions_client.py | 140 ++++++++++++++++++ 2 files changed, 165 insertions(+), 3 deletions(-) diff --git a/eval_protocol/integrations/fireworks_v1_completions_client.py b/eval_protocol/integrations/fireworks_v1_completions_client.py index c971fd03..6f72dd18 100644 --- a/eval_protocol/integrations/fireworks_v1_completions_client.py +++ b/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,6 +402,15 @@ 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 ) @@ -403,9 +418,6 @@ async def create_completion_from_prompt_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) diff --git a/tests/test_fireworks_v1_completions_client.py b/tests/test_fireworks_v1_completions_client.py index d7b863f7..8daf51b2 100644 --- a/tests/test_fireworks_v1_completions_client.py +++ b/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", From e583f8558eb90aea51288b54606ee8bbb6411a19 Mon Sep 17 00:00:00 2001 From: Anni Chen Date: Mon, 6 Jul 2026 10:09:34 -0700 Subject: [PATCH 2/4] refactor: recover per-token ids and logprobs from content[] logprobs shape MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- .../fireworks_v1_completions_client.py | 87 +++++++++---------- tests/test_fireworks_v1_completions_client.py | 83 +++++++++--------- 2 files changed, 82 insertions(+), 88 deletions(-) diff --git a/eval_protocol/integrations/fireworks_v1_completions_client.py b/eval_protocol/integrations/fireworks_v1_completions_client.py index 6f72dd18..b6bbebfb 100644 --- a/eval_protocol/integrations/fireworks_v1_completions_client.py +++ b/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) diff --git a/tests/test_fireworks_v1_completions_client.py b/tests/test_fireworks_v1_completions_client.py index 8daf51b2..e117fbe3 100644 --- a/tests/test_fireworks_v1_completions_client.py +++ b/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()) From 1f29d586f654f80090b8a9c77c295bbea1befa60 Mon Sep 17 00:00:00 2001 From: Anni Chen Date: Mon, 6 Jul 2026 10:15:14 -0700 Subject: [PATCH 3/4] fix: require sampling_logprob per content[] entry, no fallback Always read the per-token logprob from content[].sampling_logprob (the exact value the sampler drew with). Fail loud if any content[] entry lacks sampling_logprob instead of substituting the rounded content[].logprob or 0.0, so silently degraded logprobs never reach RL training. Removes the now-unused _extract_entry_logprob helper. Co-authored-by: Cursor --- .../fireworks_v1_completions_client.py | 22 +++++--------- tests/test_fireworks_v1_completions_client.py | 30 ++++++++++++++++++- 2 files changed, 37 insertions(+), 15 deletions(-) diff --git a/eval_protocol/integrations/fireworks_v1_completions_client.py b/eval_protocol/integrations/fireworks_v1_completions_client.py index b6bbebfb..4c2ce8c3 100644 --- a/eval_protocol/integrations/fireworks_v1_completions_client.py +++ b/eval_protocol/integrations/fireworks_v1_completions_client.py @@ -83,19 +83,6 @@ 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 "" @@ -437,8 +424,15 @@ async def create_completion_from_prompt_ids( f"at index {index}; cannot align per-token logprobs to token ids " "without re-encoding. Refusing to return corrupted data." ) + if entry.get("sampling_logprob") is None: + raise RuntimeError( + "Fireworks /v1/completions content[] entry is missing " + f"sampling_logprob at index {index}. The sampling logprob is the " + "exact value the sampler drew with and is required for correct " + "inference KLD; refusing to substitute the rounded logprob or 0.0." + ) completion_token_ids.append(int(entry["token_id"])) - completion_logprobs.append(_extract_entry_logprob(entry)) + completion_logprobs.append(float(entry["sampling_logprob"])) completion_text = self.decode_token_ids(token_ids=completion_token_ids) if not completion_text: diff --git a/tests/test_fireworks_v1_completions_client.py b/tests/test_fireworks_v1_completions_client.py index e117fbe3..7ab0fa92 100644 --- a/tests/test_fireworks_v1_completions_client.py +++ b/tests/test_fireworks_v1_completions_client.py @@ -174,7 +174,7 @@ def _fail_tokenizer(): "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}, + {"token_id": 26108, "sampling_logprob": -1.0, "logprob": -0.99}, ] }, } @@ -189,6 +189,34 @@ def _fail_tokenizer(): asyncio.run(client.close()) +def test_raises_when_content_entry_missing_sampling_logprob(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": [ + { + "finish_reason": "stop", + "logprobs": { + "content": [ + {"token_id": 1, "sampling_logprob": -0.1}, + {"token_id": 2, "logprob": -0.2}, + ] + }, + } + ], + }, + ) + with pytest.raises(RuntimeError, match="missing .*sampling_logprob"): + asyncio.run(client.create_completion_from_prompt_ids(prompt_token_ids=[1])) + asyncio.run(client.close()) + + def test_raises_when_no_content_logprobs(monkeypatch): client = FireworksV1CompletionsClient( model_id="test-model", From 03facac8abda4df9884e3e0f553121cba7db6085 Mon Sep 17 00:00:00 2001 From: Anni Chen Date: Mon, 6 Jul 2026 10:21:22 -0700 Subject: [PATCH 4/4] test: use pytest.approx for logprob list comparison Co-authored-by: Cursor --- tests/test_fireworks_v1_completions_client.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_fireworks_v1_completions_client.py b/tests/test_fireworks_v1_completions_client.py index 7ab0fa92..22c1fc4b 100644 --- a/tests/test_fireworks_v1_completions_client.py +++ b/tests/test_fireworks_v1_completions_client.py @@ -184,7 +184,7 @@ def _fail_tokenizer(): result = asyncio.run(client.create_completion_from_prompt_ids(prompt_token_ids=[1, 2])) 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 result["completion_logprobs"] == pytest.approx([-0.05483185, -0.0014, -1.0]) assert len(result["completion_ids"]) == len(result["completion_logprobs"]) asyncio.run(client.close())