|
| 1 | +"""Tests for the logprobs request parameter and top_logprobs response data. |
| 2 | +
|
| 3 | +The Together API accepts ``logprobs`` as an integer between 0 and 20: the |
| 4 | +number of top tokens to return log probabilities for at each generation |
| 5 | +step, instead of only the sampled token (see issue #251). When top-k |
| 6 | +logprobs are requested, each choice's ``logprobs`` part carries a |
| 7 | +``top_logprobs`` list with one ``{token: logprob}`` dict per generated |
| 8 | +token. |
| 9 | +""" |
| 10 | + |
| 11 | +import json |
| 12 | +import warnings |
| 13 | + |
| 14 | +import pytest |
| 15 | +from pydantic import ValidationError |
| 16 | + |
| 17 | +from together.types import ChatCompletionRequest, CompletionRequest |
| 18 | +from together.types.chat_completions import ChatCompletionResponse |
| 19 | + |
| 20 | + |
| 21 | +MESSAGES = [{"role": "user", "content": "Say hello."}] |
| 22 | +MODEL = "meta-llama/Llama-3.3-70B-Instruct-Turbo" |
| 23 | + |
| 24 | +# Response shape documented for the chat completions API with ``logprobs=3``: |
| 25 | +# ``top_logprobs`` is a list with one dict of the top-k alternatives per |
| 26 | +# generated token. |
| 27 | +TOP_LOGPROBS = [ |
| 28 | + {"Hello": -2.6e-06, "hello": -13.5, " Hello": -13.875}, |
| 29 | + {".": -4.8e-05, "!": -10.0625, ".\n": -11.4375}, |
| 30 | +] |
| 31 | +RESPONSE_PAYLOAD = { |
| 32 | + "id": "889ee12e7b0b3c67", |
| 33 | + "object": "chat.completion", |
| 34 | + "created": 1709240335, |
| 35 | + "model": MODEL, |
| 36 | + "choices": [ |
| 37 | + { |
| 38 | + "index": 0, |
| 39 | + "finish_reason": "eos", |
| 40 | + "logprobs": { |
| 41 | + "tokens": ["Hello", "."], |
| 42 | + "token_logprobs": [-2.6e-06, -4.8e-05], |
| 43 | + "top_logprobs": TOP_LOGPROBS, |
| 44 | + }, |
| 45 | + "message": {"role": "assistant", "content": "Hello."}, |
| 46 | + } |
| 47 | + ], |
| 48 | + "usage": {"prompt_tokens": 4, "completion_tokens": 2, "total_tokens": 6}, |
| 49 | +} |
| 50 | + |
| 51 | +# A value the API rejects, used to trigger the range error. |
| 52 | +OUT_OF_RANGE = 50 |
| 53 | +# Stand in for private prompt text, so a leak is easy to spot. |
| 54 | +SECRET_PROMPT = "unique prompt text that must not escape into an exception" |
| 55 | + |
| 56 | +RANGE_ERROR_TYPES = {"greater_than_equal", "less_than_equal"} |
| 57 | + |
| 58 | + |
| 59 | +def _assert_scoped_range_error(error: ValidationError, sent: int) -> None: |
| 60 | + """The error must name the logprobs field and carry only that value. |
| 61 | +
|
| 62 | + Scoping matters for privacy as well as for clarity. A model level |
| 63 | + validator reports the whole request as the offending input, which places |
| 64 | + the prompt or the message list inside the exception, and from there into |
| 65 | + any log line or crash reporter that serializes ``errors()``. |
| 66 | + """ |
| 67 | + details = error.errors() |
| 68 | + assert len(details) == 1 |
| 69 | + assert details[0]["loc"] == ("logprobs",) |
| 70 | + assert details[0]["type"] in RANGE_ERROR_TYPES |
| 71 | + assert details[0]["input"] == sent |
| 72 | + |
| 73 | + |
| 74 | +@pytest.mark.parametrize("logprobs", [-1, 21, 100]) |
| 75 | +def test_chat_request_rejects_out_of_range_logprobs(logprobs: int) -> None: |
| 76 | + with pytest.raises(ValidationError) as excinfo: |
| 77 | + ChatCompletionRequest(model=MODEL, messages=MESSAGES, logprobs=logprobs) |
| 78 | + |
| 79 | + _assert_scoped_range_error(excinfo.value, logprobs) |
| 80 | + |
| 81 | + |
| 82 | +@pytest.mark.parametrize("logprobs", [-1, 21, 100]) |
| 83 | +def test_completion_request_rejects_out_of_range_logprobs(logprobs: int) -> None: |
| 84 | + with pytest.raises(ValidationError) as excinfo: |
| 85 | + CompletionRequest(model=MODEL, prompt="Say hello.", logprobs=logprobs) |
| 86 | + |
| 87 | + _assert_scoped_range_error(excinfo.value, logprobs) |
| 88 | + |
| 89 | + |
| 90 | +@pytest.mark.parametrize("logprobs", [0, 1, 20]) |
| 91 | +def test_chat_request_accepts_in_range_logprobs(logprobs: int) -> None: |
| 92 | + request = ChatCompletionRequest(model=MODEL, messages=MESSAGES, logprobs=logprobs) |
| 93 | + |
| 94 | + # 0 is a valid value and must survive serialization of the payload. |
| 95 | + assert request.model_dump(exclude_none=True)["logprobs"] == logprobs |
| 96 | + |
| 97 | + |
| 98 | +@pytest.mark.parametrize("logprobs", [0, 1, 20]) |
| 99 | +def test_completion_request_accepts_in_range_logprobs(logprobs: int) -> None: |
| 100 | + request = CompletionRequest(model=MODEL, prompt="Say hello.", logprobs=logprobs) |
| 101 | + |
| 102 | + assert request.model_dump(exclude_none=True)["logprobs"] == logprobs |
| 103 | + |
| 104 | + |
| 105 | +def test_request_logprobs_defaults_to_omitted() -> None: |
| 106 | + request = ChatCompletionRequest(model=MODEL, messages=MESSAGES) |
| 107 | + |
| 108 | + assert "logprobs" not in request.model_dump(exclude_none=True) |
| 109 | + |
| 110 | + |
| 111 | +def test_chat_range_error_leaves_the_messages_out_of_the_exception() -> None: |
| 112 | + """Rejecting a bad logprobs value must not expose the conversation. |
| 113 | +
|
| 114 | + Before this check existed the request reached the server and came back as |
| 115 | + an API error that did not repeat the prompt. A client side check has to |
| 116 | + keep that property, otherwise moving validation earlier would hand user |
| 117 | + text to every logger that records the traceback. |
| 118 | + """ |
| 119 | + with pytest.raises(ValidationError) as excinfo: |
| 120 | + ChatCompletionRequest( |
| 121 | + model=MODEL, |
| 122 | + messages=[{"role": "user", "content": SECRET_PROMPT}], |
| 123 | + logprobs=OUT_OF_RANGE, |
| 124 | + ) |
| 125 | + |
| 126 | + assert SECRET_PROMPT not in str(excinfo.value) |
| 127 | + assert SECRET_PROMPT not in json.dumps(excinfo.value.errors(), default=str) |
| 128 | + |
| 129 | + |
| 130 | +def test_completion_range_error_leaves_the_prompt_out_of_the_exception() -> None: |
| 131 | + with pytest.raises(ValidationError) as excinfo: |
| 132 | + CompletionRequest(model=MODEL, prompt=SECRET_PROMPT, logprobs=OUT_OF_RANGE) |
| 133 | + |
| 134 | + assert SECRET_PROMPT not in str(excinfo.value) |
| 135 | + assert SECRET_PROMPT not in json.dumps(excinfo.value.errors(), default=str) |
| 136 | + |
| 137 | + |
| 138 | +def test_top_logprobs_survive_parsing_and_model_dump() -> None: |
| 139 | + """Top-k alternatives must round-trip through the response models. |
| 140 | +
|
| 141 | + ``top_logprobs`` is not a declared field on ``LogprobsPart``, which |
| 142 | + declares only ``tokens`` and ``token_logprobs``. It survives because the |
| 143 | + SDK's base model sets ``extra="allow"``, so the value is carried through |
| 144 | + untyped and dumped back unchanged. This test pins that pass-through: it |
| 145 | + fails if anyone later declares the field with the wrong shape, for example |
| 146 | + as a ``Dict[str, float]`` rather than a list of per-token dicts. It does |
| 147 | + not constrain what the server sends. |
| 148 | + """ |
| 149 | + response = ChatCompletionResponse(**RESPONSE_PAYLOAD) |
| 150 | + |
| 151 | + assert response.choices is not None |
| 152 | + logprobs_part = response.choices[0].logprobs |
| 153 | + assert logprobs_part is not None |
| 154 | + assert logprobs_part.top_logprobs == TOP_LOGPROBS |
| 155 | + |
| 156 | + with warnings.catch_warnings(): |
| 157 | + warnings.simplefilter("error") |
| 158 | + dumped = response.model_dump() |
| 159 | + |
| 160 | + assert dumped["choices"][0]["logprobs"]["top_logprobs"] == TOP_LOGPROBS |
0 commit comments