diff --git a/src/art/trajectories/_tokenize.py b/src/art/trajectories/_tokenize.py index 436a4814a..58fe012df 100644 --- a/src/art/trajectories/_tokenize.py +++ b/src/art/trajectories/_tokenize.py @@ -780,6 +780,9 @@ def _sampled_evidence_fingerprint( choice for choice in exchange.response.choices if choice.index == index ) choice_extra = choice.model_extra or {} + prompt = choice_extra.get("prompt_token_ids") + if prompt is None: + prompt = (exchange.response.model_extra or {}).get("prompt_token_ids") evidence = { "message": choice.message.model_dump( mode="json", @@ -795,6 +798,7 @@ def _sampled_evidence_fingerprint( }, exclude_none=True, ), + "prompt_token_ids": prompt, "token_ids": choice_extra.get("token_ids"), "logprobs": _chat_logprob_fingerprint_evidence(choice), "finish_reason": choice.finish_reason, @@ -806,9 +810,13 @@ def _sampled_evidence_fingerprint( choice for choice in exchange.response.choices if choice.index == index ) choice_extra = choice.model_extra or {} + prompt = choice_extra.get("prompt_token_ids") + if prompt is None: + prompt = (exchange.response.model_extra or {}).get("prompt_token_ids") logprobs = _dump(choice.logprobs) evidence = { "text": choice.text, + "prompt_token_ids": prompt, "token_ids": choice_extra.get("token_ids"), "logprobs": { key: logprobs[key] @@ -825,7 +833,7 @@ def _sampled_evidence_fingerprint( generation = _string_dict(generations[index]) or {} evidence = { key: generation[key] - for key in ("output_tokens", "output_indices") + for key in ("prompt_token_ids", "output_tokens", "output_indices") if key in generation } else: @@ -842,6 +850,7 @@ def _sampled_evidence_fingerprint( block.model_dump(mode="json", exclude_none=True) for block in exchange.response.content ], + "prompt_token_ids": response_extra.get("prompt_token_ids"), "token_ids": response_extra.get("token_ids"), "logprobs": response_extra.get("logprobs"), "stop_reason": exchange.response.stop_reason, @@ -865,7 +874,9 @@ def _source_key( prompt_index=prompt_index, # Internal projection may copy an exchange to isolate one choice. A # source-specific identity remains stable across those copies without - # hashing a growing request or unrelated choices. + # hashing rendered request data or unrelated choices. Captured prompt + # IDs remain part of the identity: equal output evidence can have + # different causal contexts even when response IDs are reused. evidence_fingerprint=_sampled_evidence_fingerprint( exchange, protocol=protocol, index=index ), diff --git a/tests/unit/trajectories/test_sampled_source_identity.py b/tests/unit/trajectories/test_sampled_source_identity.py new file mode 100644 index 000000000..0399d53ef --- /dev/null +++ b/tests/unit/trajectories/test_sampled_source_identity.py @@ -0,0 +1,160 @@ +from __future__ import annotations + +import pytest +from test_tokenize import ( + _chat_exchange, + _completion_exchange, + _message_exchange, + _response_exchange, +) + +import art +from art.trajectories import MessagesRequest, TrajectoryExchanges +from art.trajectories import TokenFlag as F +from art.trajectories import _tokenize as native + + +def captured(protocol): + if protocol == "chat": + return _chat_exchange([1], [2]) + if protocol == "completions": + return _completion_exchange(prompt=[1]) + if protocol == "messages": + return _message_exchange( + MessagesRequest( + model="test/model", + max_tokens=16, + messages=[{"role": "user", "content": "question"}], + ), + prompt_token_ids=[1], + token_ids=[2], + logprobs=[-0.2], + ) + return _response_exchange("response-0", 2, prompt_token_ids=[1]) + + +def change_prompt(exchange, protocol, prompt): + if protocol in {"chat", "completions"}: + exchange.response.choices[0].model_extra["prompt_token_ids"] = prompt + elif protocol == "messages": + exchange.response.model_extra["prompt_token_ids"] = prompt + else: + exchange.response.model_extra["token_generations"][0]["prompt_token_ids"] = ( + prompt + ) + + +@pytest.mark.parametrize("protocol", ["chat", "messages", "responses", "completions"]) +def test_same_response_and_output_binds_distinct_captured_prompt(protocol): + original = captured(protocol) + changed = original.model_copy(deep=True) + change_prompt(changed, protocol, [9]) + assert original.response.id == changed.response.id + assert (original.start_time, original.end_time) == ( + changed.start_time, + changed.end_time, + ) + assert native._exchange_sampled_source_key( + original + ) != native._exchange_sampled_source_key(changed) + + +@pytest.mark.parametrize("protocol", ["chat", "messages", "responses", "completions"]) +def test_identical_capture_roundtrip_retains_source_identity(protocol): + original = captured(protocol) + restored = type(original).model_validate_json(original.model_dump_json()) + assert native._exchange_sampled_source_key( + original + ) == native._exchange_sampled_source_key(restored) + + +@pytest.mark.parametrize("protocol", ["chat", "completions"]) +@pytest.mark.parametrize("choice_prompt", ["missing", "null"]) +def test_choice_without_prompt_binds_authoritative_response_fallback( + protocol, choice_prompt +): + original = captured(protocol) + extra = original.response.choices[0].model_extra + if choice_prompt == "missing": + extra.pop("prompt_token_ids") + else: + extra["prompt_token_ids"] = None + original.response.model_extra["prompt_token_ids"] = [1] + changed = original.model_copy(deep=True) + changed.response.model_extra["prompt_token_ids"] = [9] + assert native._exchange_sampled_source_key( + original + ) != native._exchange_sampled_source_key(changed) + + +@pytest.mark.parametrize("protocol", ["chat", "completions"]) +def test_selected_choice_ignores_unrelated_choices_and_shadowed_fallback(protocol): + original = captured(protocol) + original.response.choices[0].index = 7 + original.response.model_extra["prompt_token_ids"] = [900] + key = native._exchange_sampled_source_key(original) + changed = original.model_copy(deep=True) + unrelated = changed.response.choices[0].model_copy(deep=True) + unrelated.index = 3 + unrelated.model_extra["prompt_token_ids"] = [800] + changed.response.choices.insert(0, unrelated) + changed.response.model_extra["prompt_token_ids"] = [901] + assert ( + native._source_key( + changed, protocol=key.protocol, index=7, prompt_index=key.prompt_index + ) + == key + ) + unrelated.model_extra["prompt_token_ids"] = [801] + assert ( + native._source_key( + changed, protocol=key.protocol, index=7, prompt_index=key.prompt_index + ) + == key + ) + + +@pytest.mark.parametrize("multi_history", [False, True]) +def test_reused_response_identity_preserves_both_sampled_causal_prefixes(multi_history): + first = _chat_exchange([1], [2]) + second = _chat_exchange([1, 2, 3], [2], offset=1) + second.response.id = first.response.id + second.response.created = first.response.created + second.start_time = first.start_time + second.end_time = first.end_time + trajectory = art.Trajectory( + exchanges=TrajectoryExchanges(chat_completions=[first, second]) + ) + result = trajectory.tokenize(multi_history=multi_history) + histories = ( + result.histories + if isinstance(result, art.trajectories.TokenizedMultiHistoryTrajectory) + else [result] + ) + selected = { + (tuple(history.tokens[: index + 1]), history.logprobs[index]) + for history in histories + for index, flag in enumerate(history.flags) + if flag & F.SAMPLED + } + assert selected == {((1, 2), -0.2), ((1, 2, 3, 2), -0.2)} + assert ( + sum( + sum(mask) + for mask in art.trajectories.first_occurrence_masks( + histories, where=F.SAMPLED + ) + ) + == 2 + ) + restored = art.Trajectory.model_validate_json( + trajectory.model_dump_json() + ).tokenize(multi_history=multi_history) + restored_histories = ( + restored.histories + if isinstance(restored, art.trajectories.TokenizedMultiHistoryTrajectory) + else [restored] + ) + assert [(h.tokens, h.flags) for h in restored_histories] == [ + (h.tokens, h.flags) for h in histories + ]