diff --git a/CHANGELOG.md b/CHANGELOG.md index 5c34e2729..2d859fe0d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +- fix(server): accept max_completion_tokens in chat requests by @laulauraluo in #2386 + ## [0.3.36] - feat: update llama.cpp to ggml-org/llama.cpp@0c1e57098 diff --git a/llama_cpp/server/types.py b/llama_cpp/server/types.py index eda2b9280..055d39ef5 100644 --- a/llama_cpp/server/types.py +++ b/llama_cpp/server/types.py @@ -3,7 +3,7 @@ from typing import List, Optional, Union, Dict from typing_extensions import TypedDict, Literal -from pydantic import BaseModel, Field +from pydantic import AliasChoices, BaseModel, Field import llama_cpp @@ -212,6 +212,7 @@ class CreateChatCompletionRequest(BaseModel): ) # TODO: verify max_tokens: Optional[int] = Field( default=None, + validation_alias=AliasChoices("max_tokens", "max_completion_tokens"), description="The maximum number of tokens to generate. Defaults to inf", ) min_tokens: int = min_tokens_field diff --git a/tests/test_server_types.py b/tests/test_server_types.py new file mode 100644 index 000000000..bf765cefb --- /dev/null +++ b/tests/test_server_types.py @@ -0,0 +1,33 @@ +import pytest +from pydantic import ValidationError + +from llama_cpp.server.types import CreateChatCompletionRequest + + +@pytest.mark.parametrize( + ("limits", "expected"), + [ + ({}, None), + ({"max_tokens": 3}, 3), + ({"max_completion_tokens": 3}, 3), + ({"max_tokens": None}, None), + ({"max_completion_tokens": None}, None), + ({"max_tokens": 7, "max_completion_tokens": 3}, 7), + ({"max_tokens": None, "max_completion_tokens": 3}, None), + ], +) +def test_chat_completion_token_limit(limits, expected): + request = CreateChatCompletionRequest( + messages=[{"role": "user", "content": "Hello"}], **limits + ) + + assert request.max_tokens == expected + payload = request.model_dump() + assert payload["max_tokens"] == expected + assert "max_completion_tokens" not in payload + + +@pytest.mark.parametrize("field", ["max_tokens", "max_completion_tokens"]) +def test_chat_completion_token_limit_validation(field): + with pytest.raises(ValidationError): + CreateChatCompletionRequest(**{field: "invalid"})