Fix chat completion request and logprob contracts

Signed-off-by: Ritwij Aryan Parmar <ritwij.aryan.parmar@gmail.com>

Author: RitwijParmar

src/together/resources/chat/completions.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, AsyncGenerator, Dict, Iterator, List
+from typing import Any, AsyncGenerator, Dict, Iterator, List, Literal
 
 from together.abstract import api_requestor
 from together.together_response import TogetherResponse
@@ -32,6 +32,7 @@ def create(
         frequency_penalty: float | None = None,
         min_p: float | None = None,
         logit_bias: Dict[str, float] | None = None,
+        context_length_exceeded_behavior: Literal["truncate", "error"] | None = None,
         seed: int | None = None,
         stream: bool = False,
         logprobs: int | None = None,
@@ -80,6 +81,9 @@ def create(
             logit_bias (Dict[str, float], optional): A dictionary of tokens and their bias values that modify the
                 likelihood of specific tokens being sampled. Bias values must be in the range [-100, 100].
                 Defaults to None.
+            context_length_exceeded_behavior ("truncate" | "error", optional): Behavior when max_tokens exceeds the
+                model context length. "error" returns a 400, while "truncate" overrides max_tokens with the model's
+                maximum context length.
             seed (int, optional): A seed value to use for reproducibility.
             stream (bool, optional): Flag indicating whether to stream the generated completions.
                 Defaults to False.
@@ -126,6 +130,7 @@ def create(
             frequency_penalty=frequency_penalty,
             min_p=min_p,
             logit_bias=logit_bias,
+            context_length_exceeded_behavior=context_length_exceeded_behavior,
             seed=seed,
             stream=stream,
             logprobs=logprobs,
@@ -174,6 +179,7 @@ async def create(
         frequency_penalty: float | None = None,
         min_p: float | None = None,
         logit_bias: Dict[str, float] | None = None,
+        context_length_exceeded_behavior: Literal["truncate", "error"] | None = None,
         seed: int | None = None,
         stream: bool = False,
         logprobs: int | None = None,
@@ -222,6 +228,9 @@ async def create(
             logit_bias (Dict[str, float], optional): A dictionary of tokens and their bias values that modify the
                 likelihood of specific tokens being sampled. Bias values must be in the range [-100, 100].
                 Defaults to None.
+            context_length_exceeded_behavior ("truncate" | "error", optional): Behavior when max_tokens exceeds the
+                model context length. "error" returns a 400, while "truncate" overrides max_tokens with the model's
+                maximum context length.
             seed (int, optional): A seed value to use for reproducibility.
             stream (bool, optional): Flag indicating whether to stream the generated completions.
                 Defaults to False.
@@ -268,6 +277,7 @@ async def create(
             frequency_penalty=frequency_penalty,
             min_p=min_p,
             logit_bias=logit_bias,
+            context_length_exceeded_behavior=context_length_exceeded_behavior,
             seed=seed,
             stream=stream,
             logprobs=logprobs,

src/together/types/chat_completions.py

@@ -2,7 +2,7 @@
 
 import warnings
 from enum import Enum
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Literal
 
 from pydantic import model_validator
 from typing_extensions import Self
@@ -132,6 +132,8 @@ class ChatCompletionRequest(BaseModel):
     frequency_penalty: float | None = None
     min_p: float | None = None
     logit_bias: Dict[str, float] | None = None
+    # behavior when max_tokens exceeds the model context length
+    context_length_exceeded_behavior: Literal["truncate", "error"] | None = None
     seed: int | None = None
     # stream SSE token chunks
     stream: bool = False

src/together/types/common.py

@@ -42,6 +42,8 @@ class LogprobsPart(BaseModel):
     tokens: List[str | None] | None = None
     # token logprob list
     token_logprobs: List[float | None] | None = None
+    # top-k logprobs per token
+    top_logprobs: List[Dict[str, float]] | None = None
 
 
 class PromptPart(BaseModel):

tests/unit/test_chat_completion_contract.py

@@ -0,0 +1,59 @@
+import inspect
+
+from together.resources.chat.completions import AsyncChatCompletions, ChatCompletions
+from together.types.chat_completions import (
+    ChatCompletionRequest,
+    ChatCompletionResponse,
+)
+from together.types.common import LogprobsPart
+
+
+def test_chat_completion_create_exposes_context_length_behavior() -> None:
+    sync_signature = inspect.signature(ChatCompletions.create)
+    async_signature = inspect.signature(AsyncChatCompletions.create)
+
+    assert "context_length_exceeded_behavior" in sync_signature.parameters
+    assert "context_length_exceeded_behavior" in async_signature.parameters
+
+
+def test_chat_completion_request_serializes_context_length_behavior() -> None:
+    request = ChatCompletionRequest(
+        model="meta-llama/Llama-3.3-70B-Instruct-Turbo",
+        messages=[{"role": "user", "content": "Hello"}],
+        context_length_exceeded_behavior="truncate",
+    )
+
+    assert (
+        request.model_dump(exclude_none=True)["context_length_exceeded_behavior"]
+        == "truncate"
+    )
+
+
+def test_logprobs_part_models_top_logprobs_as_list_per_token() -> None:
+    assert "top_logprobs" in LogprobsPart.model_fields
+
+    response = ChatCompletionResponse(
+        choices=[
+            {
+                "logprobs": {
+                    "tokens": ["Hello", "."],
+                    "token_logprobs": [-0.1, -0.2],
+                    "top_logprobs": [
+                        {"Hello": -0.1, "Hi": -1.4},
+                        {".": -0.2, "!": -1.7},
+                    ],
+                }
+            }
+        ]
+    )
+
+    assert response.choices is not None
+    assert response.choices[0].logprobs is not None
+    assert response.choices[0].logprobs.top_logprobs == [
+        {"Hello": -0.1, "Hi": -1.4},
+        {".": -0.2, "!": -1.7},
+    ]
+    assert response.model_dump()["choices"][0]["logprobs"]["top_logprobs"] == [
+        {"Hello": -0.1, "Hi": -1.4},
+        {".": -0.2, "!": -1.7},
+    ]