diff --git a/eval_protocol/integrations/fireworks_v1_completions_client.py b/eval_protocol/integrations/fireworks_v1_completions_client.py index c971fd03..4c2ce8c3 100644 --- a/eval_protocol/integrations/fireworks_v1_completions_client.py +++ b/eval_protocol/integrations/fireworks_v1_completions_client.py @@ -393,35 +393,50 @@ 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 [] - ) 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 ) + # -- 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") + content_entries = choice_logprobs.get("content") if isinstance(choice_logprobs, dict) else None + if not content_entries: + raise RuntimeError( + "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." + ) + 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(float(entry["sampling_logprob"])) + 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] = [] - 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] # -- Build message via parser or raw -------------------------------- if self.tool_call_parser is not None: diff --git a/tests/test_fireworks_v1_completions_client.py b/tests/test_fireworks_v1_completions_client.py index d7b863f7..22c1fc4b 100644 --- a/tests/test_fireworks_v1_completions_client.py +++ b/tests/test_fireworks_v1_completions_client.py @@ -124,6 +124,169 @@ 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_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: "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": [ + { + "finish_reason": "stop", + "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, "sampling_logprob": -1.0, "logprob": -0.99}, + ] + }, + } + ], + }, + ) + 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"] == pytest.approx([-0.05483185, -0.0014, -1.0]) + assert len(result["completion_ids"]) == len(result["completion_logprobs"]) + 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", + 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 content\\[\\] logprobs entries"): + asyncio.run(client.create_completion_from_prompt_ids(prompt_token_ids=[1, 2])) + asyncio.run(client.close()) + + +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", + ) + _install_fake_completion( + client, + monkeypatch, + { + "choices": [ + { + "finish_reason": "stop", + "token_ids": [1, 2, 3], + "logprobs": { + "token_ids": [1, 2, 3], + "token_logprobs": [-0.1, -0.2, -0.3], + }, + } + ], + }, + ) + 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_when_content_entry_missing_token_id(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}, + {"sampling_logprob": -0.2}, + ] + }, + } + ], + }, + ) + with pytest.raises(RuntimeError, match="missing token_id"): + 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",