Validate logprobs range and pin top-k logprobs round-trip

logprobs must be 0 to 20, but out of range values were only caught by the
API. Reject them in the request validator, document the range in the
create() docstrings, and add a test that top_logprobs survives
model_dump() (#251, #443).

Addresses #251.

Author: Ranoobaba

src/together/resources/chat/completions.py

@@ -83,7 +83,9 @@ def create(
             seed (int, optional): A seed value to use for reproducibility.
             stream (bool, optional): Flag indicating whether to stream the generated completions.
                 Defaults to False.
-            logprobs (int, optional): Number of top-k logprobs to return
+            logprobs (int, optional): Number of top tokens to return log probabilities for
+                    at each generation step, instead of only the sampled token.
+                    Must be in the range [0, 20].
                 Defaults to None.
             echo (bool, optional): Echo prompt in output. Can be used with logprobs to return prompt logprobs.
                 Defaults to None.
@@ -225,7 +227,9 @@ async def create(
             seed (int, optional): A seed value to use for reproducibility.
             stream (bool, optional): Flag indicating whether to stream the generated completions.
                 Defaults to False.
-            logprobs (int, optional): Number of top-k logprobs to return
+            logprobs (int, optional): Number of top tokens to return log probabilities for
+                    at each generation step, instead of only the sampled token.
+                    Must be in the range [0, 20].
                 Defaults to None.
             echo (bool, optional): Echo prompt in output. Can be used with logprobs to return prompt logprobs.
                 Defaults to None.

src/together/resources/completions.py

@@ -79,7 +79,9 @@ def create(
             seed (int, optional): Seed value for reproducibility.
             stream (bool, optional): Flag indicating whether to stream the generated completions.
                 Defaults to False.
-            logprobs (int, optional): Number of top-k logprobs to return
+            logprobs (int, optional): Number of top tokens to return log probabilities for
+                    at each generation step, instead of only the sampled token.
+                    Must be in the range [0, 20].
                 Defaults to None.
             echo (bool, optional): Echo prompt in output. Can be used with logprobs to return prompt logprobs.
                 Defaults to None.
@@ -203,7 +205,9 @@ async def create(
             seed (int, optional): Seed value for reproducibility.
             stream (bool, optional): Flag indicating whether to stream the generated completions.
                 Defaults to False.
-            logprobs (int, optional): Number of top-k logprobs to return
+            logprobs (int, optional): Number of top tokens to return log probabilities for
+                    at each generation step, instead of only the sampled token.
+                    Must be in the range [0, 20].
                 Defaults to None.
             echo (bool, optional): Echo prompt in output. Can be used with logprobs to return prompt logprobs.
                 Defaults to None.

src/together/types/chat_completions.py

@@ -150,13 +150,20 @@ class ChatCompletionRequest(BaseModel):
     tool_choice: ToolChoice | ToolChoiceEnum | None = None
 
     # Raise warning if repetition_penalty is used with presence_penalty or frequency_penalty
+    # and raise an error if logprobs is outside the range supported by the API.
     @model_validator(mode="after")
     def verify_parameters(self) -> Self:
         if self.repetition_penalty:
             if self.presence_penalty or self.frequency_penalty:
                 warnings.warn(
                     "repetition_penalty is not advisable to be used alongside presence_penalty or frequency_penalty"
                 )
+        if self.logprobs is not None and not 0 <= self.logprobs <= 20:
+            raise ValueError(
+                f"logprobs must be an integer between 0 and 20, got {self.logprobs}. "
+                "It sets the number of top tokens to return log probabilities for "
+                "at each generation step, instead of only the sampled token."
+            )
         return self
 
 

src/together/types/completions.py

@@ -49,13 +49,20 @@ class CompletionRequest(BaseModel):
     safety_model: str | None = None
 
     # Raise warning if repetition_penalty is used with presence_penalty or frequency_penalty
+    # and raise an error if logprobs is outside the range supported by the API.
     @model_validator(mode="after")
     def verify_parameters(self) -> Self:
         if self.repetition_penalty:
             if self.presence_penalty or self.frequency_penalty:
                 warnings.warn(
                     "repetition_penalty is not advisable to be used alongside presence_penalty or frequency_penalty"
                 )
+        if self.logprobs is not None and not 0 <= self.logprobs <= 20:
+            raise ValueError(
+                f"logprobs must be an integer between 0 and 20, got {self.logprobs}. "
+                "It sets the number of top tokens to return log probabilities for "
+                "at each generation step, instead of only the sampled token."
+            )
         return self
 
 

tests/unit/test_logprobs.py

@@ -0,0 +1,103 @@
+"""Tests for the logprobs request parameter and top_logprobs response data.
+
+The Together API accepts ``logprobs`` as an integer between 0 and 20: the
+number of top tokens to return log probabilities for at each generation
+step, instead of only the sampled token (see issue #251). When top-k
+logprobs are requested, each choice's ``logprobs`` part carries a
+``top_logprobs`` list with one ``{token: logprob}`` dict per generated
+token.
+"""
+
+import warnings
+
+import pytest
+from pydantic import ValidationError
+
+from together.types import ChatCompletionRequest, CompletionRequest
+from together.types.chat_completions import ChatCompletionResponse
+
+
+MESSAGES = [{"role": "user", "content": "Say hello."}]
+MODEL = "meta-llama/Llama-3.3-70B-Instruct-Turbo"
+
+# Real response shape captured from the chat completions API with
+# ``logprobs=3`` (see issues #251 and #443): ``top_logprobs`` is a list
+# with one dict of the top-k alternatives per generated token.
+TOP_LOGPROBS = [
+    {"Hello": -2.6e-06, "hello": -13.5, " Hello": -13.875},
+    {".": -4.8e-05, "!": -10.0625, ".\n": -11.4375},
+]
+RESPONSE_PAYLOAD = {
+    "id": "889ee12e7b0b3c67",
+    "object": "chat.completion",
+    "created": 1709240335,
+    "model": MODEL,
+    "choices": [
+        {
+            "index": 0,
+            "finish_reason": "eos",
+            "logprobs": {
+                "tokens": ["Hello", "."],
+                "token_logprobs": [-2.6e-06, -4.8e-05],
+                "top_logprobs": TOP_LOGPROBS,
+            },
+            "message": {"role": "assistant", "content": "Hello."},
+        }
+    ],
+    "usage": {"prompt_tokens": 4, "completion_tokens": 2, "total_tokens": 6},
+}
+
+
+@pytest.mark.parametrize("logprobs", [-1, 21, 100])
+def test_chat_request_rejects_out_of_range_logprobs(logprobs: int) -> None:
+    with pytest.raises(ValidationError, match="between 0 and 20"):
+        ChatCompletionRequest(model=MODEL, messages=MESSAGES, logprobs=logprobs)
+
+
+@pytest.mark.parametrize("logprobs", [-1, 21, 100])
+def test_completion_request_rejects_out_of_range_logprobs(logprobs: int) -> None:
+    with pytest.raises(ValidationError, match="between 0 and 20"):
+        CompletionRequest(model=MODEL, prompt="Say hello.", logprobs=logprobs)
+
+
+@pytest.mark.parametrize("logprobs", [0, 1, 20])
+def test_chat_request_accepts_in_range_logprobs(logprobs: int) -> None:
+    request = ChatCompletionRequest(model=MODEL, messages=MESSAGES, logprobs=logprobs)
+
+    # 0 is a valid value and must survive serialization of the payload.
+    assert request.model_dump(exclude_none=True)["logprobs"] == logprobs
+
+
+@pytest.mark.parametrize("logprobs", [0, 1, 20])
+def test_completion_request_accepts_in_range_logprobs(logprobs: int) -> None:
+    request = CompletionRequest(model=MODEL, prompt="Say hello.", logprobs=logprobs)
+
+    assert request.model_dump(exclude_none=True)["logprobs"] == logprobs
+
+
+def test_request_logprobs_defaults_to_omitted() -> None:
+    request = ChatCompletionRequest(model=MODEL, messages=MESSAGES)
+
+    assert "logprobs" not in request.model_dump(exclude_none=True)
+
+
+def test_top_logprobs_survive_parsing_and_model_dump() -> None:
+    """Top-k alternatives must round-trip through the response models.
+
+    Guards against the mistyping reported in issue #443, where
+    ``top_logprobs`` declared as ``Dict[str, float]`` (instead of a list
+    of per-token dicts) made ``model_dump()`` emit
+    ``PydanticSerializationUnexpectedValue`` warnings.
+    """
+    response = ChatCompletionResponse(**RESPONSE_PAYLOAD)
+
+    assert response.choices is not None
+    logprobs_part = response.choices[0].logprobs
+    assert logprobs_part is not None
+    assert logprobs_part.top_logprobs == TOP_LOGPROBS
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+        dumped = response.model_dump()
+
+    assert dumped["choices"][0]["logprobs"]["top_logprobs"] == TOP_LOGPROBS