diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py index 3e2acc854..94b3e10ef 100644 --- a/python/freetoken/server/api_server.py +++ b/python/freetoken/server/api_server.py @@ -33,7 +33,7 @@ init_logger, load_generation_sampling, ) -from pydantic import BaseModel +from pydantic import BaseModel, Field from .args import ServerArgs from .anthropic_api import register_anthropic_routes @@ -494,7 +494,7 @@ class CacheRebuildRequest(BaseModel): # is deferred (needs the drain-gate machinery); constraining the Literal makes an # unsupported value fail fast with a 422 at the API layer instead of a generic 503. mode: Literal["if_idle"] = "if_idle" - timeout: float = 300.0 + timeout: float = Field(default=300.0, gt=0, le=3600) async def dispatch_rebuild( diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py index be05d908a..a60119ecf 100644 --- a/python/freetoken/server/generation.py +++ b/python/freetoken/server/generation.py @@ -15,6 +15,7 @@ import asyncio import json +import math import time from collections.abc import AsyncIterator from dataclasses import dataclass, field @@ -177,12 +178,21 @@ def pick(value, key, framework): # non-positive value is a client error. if max_tokens is not None and max_tokens < 1: raise ValueError(f"max_tokens must be at least 1, got {max_tokens}") + resolved_temperature = pick(temperature, "temperature", 0.0) + if not math.isfinite(resolved_temperature) or resolved_temperature < 0: + raise ValueError(f"temperature must be a finite number >= 0, got {resolved_temperature}") + resolved_top_p = pick(top_p, "top_p", 1.0) + if not math.isfinite(resolved_top_p) or not 0 < resolved_top_p <= 1: + raise ValueError(f"top_p must be in (0, 1], got {resolved_top_p}") + resolved_top_k = pick(top_k, "top_k", -1) + if resolved_top_k != -1 and resolved_top_k < 1: + raise ValueError(f"top_k must be -1 (disabled) or >= 1, got {resolved_top_k}") return SamplingParams( ignore_eos=ignore_eos, max_tokens=DEFAULT_MAX_OUTPUT_TOKENS if max_tokens is None else max_tokens, - temperature=pick(temperature, "temperature", 0.0), - top_k=pick(top_k, "top_k", -1), - top_p=pick(top_p, "top_p", 1.0), + temperature=resolved_temperature, + top_k=resolved_top_k, + top_p=resolved_top_p, stop_strs=[s for s in stop_list if s], # drop empty strings (would match everything) ) diff --git a/tests/server/test_openai_api.py b/tests/server/test_openai_api.py index facd469b3..f6808a7d3 100644 --- a/tests/server/test_openai_api.py +++ b/tests/server/test_openai_api.py @@ -676,3 +676,66 @@ def test_minimax_http_non_stream_forces_implicit_reasoning_without_request_knob( message = response["choices"][0]["message"] assert message["reasoning_content"] == "private thought" assert message["content"] == "visible answer" + + +import pytest # noqa: E402 + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("temperature", -0.5), + ("top_p", 0.0), + ("top_k", 0), + ], +) +def test_chat_completion_rejects_invalid_sampling(field, value): + app = FastAPI() + state = FakeState([]) + + @app.post("/v1/chat/completions") + async def chat_completion(req: ChatCompletionRequest): + return await handle_chat_completion(req, request=None, state=state, model_sampling={}) + + body = { + "model": "unit-model", + "messages": [{"role": "user", "content": "hi"}], + "stream": True, + field: value, + } + response = TestClient(app).post("/v1/chat/completions", json=body) + + assert response.status_code == 400 + error = response.json()["error"] + assert error["type"] == "invalid_request_error" + assert field in error["message"] + assert state.sent is None + + +@pytest.mark.parametrize( + "sampling", + [ + {"temperature": 0.0, "top_p": 1.0, "top_k": -1}, + {"temperature": 0.0, "top_p": 1.0, "top_k": 1}, + ], +) +def test_chat_completion_accepts_sampling_boundaries(sampling): + app = FastAPI() + state = FakeState([]) + + @app.post("/v1/chat/completions") + async def chat_completion(req: ChatCompletionRequest): + return await handle_chat_completion(req, request=None, state=state, model_sampling={}) + + response = TestClient(app).post( + "/v1/chat/completions", + json={ + "model": "unit-model", + "messages": [{"role": "user", "content": "hi"}], + "stream": True, + **sampling, + }, + ) + + assert response.status_code == 200 + assert state.sent is not None diff --git a/tests/server/test_rebuild_maintenance.py b/tests/server/test_rebuild_maintenance.py index a1d10a7e5..0df3d6bfe 100644 --- a/tests/server/test_rebuild_maintenance.py +++ b/tests/server/test_rebuild_maintenance.py @@ -327,3 +327,18 @@ def test_cache_rebuild_request_rejects_unknown_mode(): assert CacheRebuildRequest(mode="if_idle").mode == "if_idle" with pytest.raises(ValidationError): CacheRebuildRequest(mode="drain") + + +from pydantic import ValidationError # noqa: E402 + +from freetoken.server.api_server import CacheRebuildRequest # noqa: E402 + + +def test_rebuild_timeout_bounds(): + import pytest + + for bad in (0, -5, 3601, float("nan")): + with pytest.raises(ValidationError): + CacheRebuildRequest(timeout=bad) + assert CacheRebuildRequest(timeout=3600).timeout == 3600 + assert CacheRebuildRequest().timeout == 300.0