diff --git a/python/freetoken/server/anthropic_api.py b/python/freetoken/server/anthropic_api.py index c941d98b7..18911673f 100644 --- a/python/freetoken/server/anthropic_api.py +++ b/python/freetoken/server/anthropic_api.py @@ -33,6 +33,7 @@ AnthropicUsage, ) from .generation import ( + DEFAULT_MAX_OUTPUT_TOKENS, KEEPALIVE, ContentDelta, GenDone, @@ -116,6 +117,9 @@ async def handle_anthropic_messages( spec = convert_anthropic_to_genspec( req, model_sampling, reasoning_parser=getattr(state.config, "reasoning_parser", None), + default_max_tokens=( + getattr(state.config, "max_output_tokens", None) or DEFAULT_MAX_OUTPUT_TOKENS + ), ) uid = await submit_generation(spec, state) except ValueError as exc: @@ -303,6 +307,7 @@ def convert_anthropic_to_genspec( req: AnthropicMessagesRequest, model_sampling: dict[str, Any], reasoning_parser: str | None = None, + default_max_tokens: int = DEFAULT_MAX_OUTPUT_TOKENS, ) -> GenSpec: messages, template_tools, parser_tools, ctk = convert_anthropic_prompt( req, reasoning_parser=reasoning_parser @@ -317,6 +322,7 @@ def convert_anthropic_to_genspec( ignore_eos=False, model_sampling=model_sampling, stop=req.stop_sequences, + default_max_tokens=default_max_tokens, ), chat_template_kwargs=ctk, template_tools=template_tools, diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py index be05d908a..28f85d87a 100644 --- a/python/freetoken/server/generation.py +++ b/python/freetoken/server/generation.py @@ -163,6 +163,7 @@ def resolve_sampling( ignore_eos: bool, model_sampling: dict[str, Any], stop: str | list[str] | None = None, + default_max_tokens: int = DEFAULT_MAX_OUTPUT_TOKENS, ) -> SamplingParams: """Map a protocol's sampling fields onto the engine's neutral SamplingParams, filling unspecified fields from the checkpoint's recommended defaults.""" @@ -179,7 +180,7 @@ def pick(value, key, framework): raise ValueError(f"max_tokens must be at least 1, got {max_tokens}") return SamplingParams( ignore_eos=ignore_eos, - max_tokens=DEFAULT_MAX_OUTPUT_TOKENS if max_tokens is None else max_tokens, + max_tokens=default_max_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), diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py index b4becd263..a76898abd 100644 --- a/python/freetoken/server/openai_api.py +++ b/python/freetoken/server/openai_api.py @@ -23,6 +23,7 @@ from .function_call_parser import ToolCallItem from .request_logger import log_request from .generation import ( + DEFAULT_MAX_OUTPUT_TOKENS, ContentDelta, GenDone, GenerationError, @@ -58,6 +59,7 @@ def _thinking_type(req: Any) -> str | None: def chat_request_to_genspec( req: ChatCompletionRequest, model_sampling: dict[str, Any], + default_max_tokens: int = DEFAULT_MAX_OUTPUT_TOKENS, ) -> GenSpec: """OpenAI ChatCompletionRequest -> GenSpec (the OpenAI 'to_sampling_params').""" from .model_meta import effort_toggle_kwargs @@ -76,6 +78,7 @@ def chat_request_to_genspec( ignore_eos=req.ignore_eos, model_sampling=model_sampling, stop=req.stop, + default_max_tokens=default_max_tokens, ), chat_template_kwargs=ctk, template_tools=_tools_for_template(req), @@ -179,7 +182,10 @@ async def handle_chat_completion( ) try: - spec = chat_request_to_genspec(req, model_sampling) + default_max_tokens = ( + getattr(state.config, "max_output_tokens", None) or DEFAULT_MAX_OUTPUT_TOKENS + ) + spec = chat_request_to_genspec(req, model_sampling, default_max_tokens=default_max_tokens) except ValueError as exc: return create_error_response(str(exc)) @@ -386,7 +392,10 @@ async def handle_completion( if unsupported is not None: return create_error_response(unsupported) try: # surfaces an out-of-range max_tokens as a 400 rather than a 500 from the worker - _resolve_sampling(req, model_sampling) + default_max_tokens = ( + getattr(state.config, "max_output_tokens", None) or DEFAULT_MAX_OUTPUT_TOKENS + ) + _resolve_sampling(req, model_sampling, default_max_tokens=default_max_tokens) except ValueError as exc: return create_error_response(str(exc), param="max_tokens") @@ -397,7 +406,9 @@ async def handle_completion( return create_error_response("Streaming completions only support a single text prompt") uid = state.new_user() await state.send_one( - TokenizeMsg(uid=uid, text=prompts[0], sampling_params=_resolve_sampling(req, model_sampling)) + TokenizeMsg(uid=uid, text=prompts[0], sampling_params=_resolve_sampling( + req, model_sampling, default_max_tokens=default_max_tokens + )) ) chunks = stream_completion_chunks(uid, req, state) if request is not None: @@ -410,7 +421,15 @@ async def handle_completion( cached_tokens = 0 for index, prompt in enumerate(prompts): uid = state.new_user() - await state.send_one(TokenizeMsg(uid=uid, text=prompt, sampling_params=_resolve_sampling(req, model_sampling))) + await state.send_one( + TokenizeMsg( + uid=uid, + text=prompt, + sampling_params=_resolve_sampling( + req, model_sampling, default_max_tokens=default_max_tokens + ), + ) + ) text = "" finish_reason = "stop" async for ack in state.wait_for_ack(uid): @@ -517,6 +536,7 @@ def create_error_response( def _resolve_sampling( req: ChatCompletionRequest | CompletionRequest, model_sampling: dict[str, Any], + default_max_tokens: int = DEFAULT_MAX_OUTPUT_TOKENS, ) -> SamplingParams: return resolve_sampling( temperature=req.temperature, @@ -526,6 +546,7 @@ def _resolve_sampling( ignore_eos=req.ignore_eos, model_sampling=model_sampling, stop=req.stop, + default_max_tokens=default_max_tokens, ) diff --git a/tests/server/test_openai_api.py b/tests/server/test_openai_api.py index facd469b3..58d15d7c1 100644 --- a/tests/server/test_openai_api.py +++ b/tests/server/test_openai_api.py @@ -391,31 +391,41 @@ def test_completion_forwards_length_finish_reason(): assert response["choices"][0]["finish_reason"] == "length" -def test_omitted_max_tokens_defaults_to_hardcoded_32k(): +def test_omitted_max_tokens_honors_server_default(): from freetoken.server.generation import DEFAULT_MAX_OUTPUT_TOKENS chat_state = FakeState([UserReply(uid=42, incremental_output="hi", finished=True)]) + chat_state.config.max_output_tokens = 4096 run(handle_chat_completion( ChatCompletionRequest(model="m", messages=[{"role": "user", "content": "hi"}]), request=None, state=chat_state, model_sampling={}, )) - assert chat_state.sent.sampling_params.max_tokens == DEFAULT_MAX_OUTPUT_TOKENS + assert chat_state.sent.sampling_params.max_tokens == 4096 cmpl_state = FakeState([UserReply(uid=42, incremental_output="hi", finished=True)]) + cmpl_state.config.max_output_tokens = 4096 run(handle_completion( CompletionRequest(model="m", prompt="hi"), request=None, state=cmpl_state, model_sampling={}, )) - assert cmpl_state.sent.sampling_params.max_tokens == DEFAULT_MAX_OUTPUT_TOKENS + assert cmpl_state.sent.sampling_params.max_tokens == 4096 # explicit value wins exp_state = FakeState([UserReply(uid=42, incremental_output="hi", finished=True)]) + exp_state.config.max_output_tokens = 4096 run(handle_completion( CompletionRequest(model="m", prompt="hi", max_tokens=50), request=None, state=exp_state, model_sampling={}, )) assert exp_state.sent.sampling_params.max_tokens == 50 + fallback_state = FakeState([UserReply(uid=42, incremental_output="hi", finished=True)]) + run(handle_chat_completion( + ChatCompletionRequest(model="m", messages=[{"role": "user", "content": "hi"}]), + request=None, state=fallback_state, model_sampling={}, + )) + assert fallback_state.sent.sampling_params.max_tokens == DEFAULT_MAX_OUTPUT_TOKENS + def test_models_route_returns_served_model_name(): state = FakeState([])