diff --git a/Cargo.lock b/Cargo.lock index 93ee7e07d4af..f72ab6dbed98 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1574,7 +1574,7 @@ version = "3.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "faf9468729b8cbcea668e36183cb69d317348c2e08e994829fb56ebfdfbaac34" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.48.0", ] [[package]] @@ -2901,9 +2901,8 @@ dependencies = [ [[package]] name = "dynamo-parsers-v2" -version = "0.7.21" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f43e47b6a68353cc04d4dfe19c69414a5884c9963b72c68937d019998c5ef72e" +version = "0.10.0" +source = "git+https://github.com/ai-dynamo/frontend-crates.git?rev=ef8f401393a730cae8ac5c96fe1d564dcf48a017#ef8f401393a730cae8ac5c96fe1d564dcf48a017" dependencies = [ "anyhow", "openai-harmony", @@ -2936,9 +2935,8 @@ dependencies = [ [[package]] name = "dynamo-renderer" -version = "8.0.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9dc3c1f04389ad78e3f9dad14360a191c27f3e3f574aba108b9986f5592cfea3" +version = "8.1.0" +source = "git+https://github.com/ai-dynamo/frontend-crates.git?rev=ef8f401393a730cae8ac5c96fe1d564dcf48a017#ef8f401393a730cae8ac5c96fe1d564dcf48a017" dependencies = [ "anyhow", "chrono", @@ -7735,8 +7733,8 @@ version = "0.13.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "be769465445e8c1474e9c5dac2018218498557af32d9ed057325ec9a41ae81bf" dependencies = [ - "heck 0.5.0", - "itertools 0.14.0", + "heck 0.4.1", + "itertools 0.10.5", "log", "multimap", "once_cell", @@ -7755,8 +7753,8 @@ version = "0.14.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "343d3bd7056eda839b03204e68deff7d1b13aba7af2b2fd16890697274262ee7" dependencies = [ - "heck 0.5.0", - "itertools 0.14.0", + "heck 0.4.1", + "itertools 0.10.5", "log", "multimap", "petgraph 0.8.3", @@ -7777,7 +7775,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a56d757972c98b346a9b766e3f02746cde6dd1cd1d1d563472929fdd74bec4d" dependencies = [ "anyhow", - "itertools 0.14.0", + "itertools 0.10.5", "proc-macro2", "quote", "syn 2.0.117", @@ -7790,7 +7788,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "27c6023962132f4b30eb4c172c91ce92d933da334c59c23cddee82358ddafb0b" dependencies = [ "anyhow", - "itertools 0.14.0", + "itertools 0.10.5", "proc-macro2", "quote", "syn 2.0.117", @@ -11512,7 +11510,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.48.0", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 43308867dc74..0f329b2c7a50 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -66,7 +66,10 @@ dynamo-runtime = { path = "lib/runtime", version = "1.6.0" } dynamo-llm = { path = "lib/llm", version = "1.6.0" } dynamo-rl = { path = "lib/rl", version = "1.6.0" } dynamo-tokenizers = { version = "=3.1.0" } -dynamo-renderer = { version = "=8.0.1" } +# Beam release pin: renderer 8.1.0 (the Beam `builtin_tools` rendering view) comes from the +# frontend-crates `release/1.6.0-beam-1.0-501b-dev.1` line and is not published to crates.io. +# Same revision as `dynamo-parsers-v2` below, so one immutable commit covers both crates. +dynamo-renderer = { git = "https://github.com/ai-dynamo/frontend-crates.git", rev = "ef8f401393a730cae8ac5c96fe1d564dcf48a017", version = "=8.1.0" } dynamo-tokens = { path = "lib/tokens", version = "1.6.0" } dynamo-truthy = { path = "lib/truthy", version = "1.6.0" } dynamo-data-gen = { path = "lib/data-gen", version = "1.6.0" } @@ -81,7 +84,9 @@ dynamo-parsers = { version = "=10.0.2" } # web: https://crates.io/crates/dynamo-parsers-v2/versions # sparse index: curl https://index.crates.io/dy/na/dynamo-parsers-v2 (one JSON line per release; read the "vers" field) # Native count-before-string early streaming still needs a frontend-crates fix. -dynamo-parsers-v2 = { version = "=0.7.21" } +# Beam release pin: parsers-v2 0.10.0 (the Beam unified family) comes from the frontend-crates +# `release/1.6.0-beam-1.0-501b-dev.1` line and is not published to crates.io. +dynamo-parsers-v2 = { git = "https://github.com/ai-dynamo/frontend-crates.git", rev = "ef8f401393a730cae8ac5c96fe1d564dcf48a017", version = "=0.10.0" } dynamo-backend-common = { path = "lib/backend-common", version = "1.6.0" } dynamo-sidecar-common = { path = "lib/sidecar/common", version = "1.6.0" } diff --git a/components/src/dynamo/frontend/prepost.py b/components/src/dynamo/frontend/prepost.py index d4004f18095a..c150b1c6da6d 100644 --- a/components/src/dynamo/frontend/prepost.py +++ b/components/src/dynamo/frontend/prepost.py @@ -14,7 +14,13 @@ from functools import lru_cache, partial from typing import TYPE_CHECKING, Any, Protocol, TypeGuard, cast -from vllm.entrypoints.chat_utils import make_tool_call_id +try: + from vllm.entrypoints.chat_utils import make_tool_call_id +except ImportError: + # Newer vLLM moved chat_utils under vllm.renderers (renderers refactor, + # upstream main after 2026-10). Remove this fallback once every supported + # vLLM pin ships vllm.renderers.chat_utils. + from vllm.renderers.chat_utils.tool_call import make_tool_call_id from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionNamedFunction, ChatCompletionNamedToolChoiceParam, diff --git a/components/src/dynamo/frontend/vllm_processor.py b/components/src/dynamo/frontend/vllm_processor.py index 1e2a75f20534..068e6df27eea 100644 --- a/components/src/dynamo/frontend/vllm_processor.py +++ b/components/src/dynamo/frontend/vllm_processor.py @@ -18,7 +18,15 @@ from msgspec.structs import replace as msgspec_replace from vllm.config import CacheConfig, LoadConfig, ModelConfig, VllmConfig -from vllm.entrypoints.chat_utils import load_chat_template + +try: + from vllm.entrypoints.chat_utils import load_chat_template +except ImportError: + # Newer vLLM moved chat_utils under vllm.renderers (renderers refactor, + # upstream main after 2026-10). Remove this fallback once every supported + # vLLM pin ships vllm.renderers.chat_utils. + from vllm.renderers.chat_utils.chat_template import load_chat_template + from vllm.exceptions import VLLMClientError from vllm.reasoning import ReasoningParser, ReasoningParserManager from vllm.sampling_params import RequestOutputKind, SamplingParams diff --git a/components/src/dynamo/vllm/embedding_worker_processes.py b/components/src/dynamo/vllm/embedding_worker_processes.py index 06d0c2343d62..7947759a0b26 100644 --- a/components/src/dynamo/vllm/embedding_worker_processes.py +++ b/components/src/dynamo/vllm/embedding_worker_processes.py @@ -13,6 +13,7 @@ from __future__ import annotations import copy +import importlib import json import logging import os @@ -29,9 +30,43 @@ from vllm.config import VllmConfig from vllm.usage.usage_lib import UsageContext from vllm.v1.engine.async_llm import AsyncLLM -from vllm.v1.engine.utils import get_engine_zmq_addresses, launch_core_engines from vllm.v1.executor import Executor +# Resolved on first use, not at import: vLLM renamed get_engine_zmq_addresses in +# vllm-project/vllm#54113, and this module is imported by every dynamo.vllm worker +# while only --embedding-worker-processes > 1 calls into it. Module attributes so +# tests can substitute them. +get_engine_zmq_addresses: Callable[..., Any] | None = None +launch_core_engines: Callable[..., Any] | None = None + + +def _resolve_engine_utils() -> tuple[Callable[..., Any], Callable[..., Any]]: + """Bind the vLLM engine-launch helpers this pool needs, failing with a clear message. + + Returns ``(get_engine_zmq_addresses, launch_core_engines)`` so callers hold + non-optional references; the module attributes are kept in sync for tests. + """ + global get_engine_zmq_addresses, launch_core_engines + engine_utils = importlib.import_module("vllm.v1.engine.utils") + + if launch_core_engines is None: + launch_core_engines = engine_utils.launch_core_engines + if get_engine_zmq_addresses is None: + resolved = getattr(engine_utils, "get_engine_zmq_addresses", None) + if resolved is None: + # vllm#54113: the supervisor now binds the frontend listeners up front and + # children adopt them; this pool still lets children bind the addresses + # themselves, so refuse rather than hand out addresses that are already bound. + raise RuntimeError( + "--embedding-worker-processes > 1 is not yet supported on this vLLM: " + f"vLLM {vllm.__version__} has no get_engine_zmq_addresses (it binds " + "frontend listeners in the supervisor since vllm-project/vllm#54113). " + "Run with --embedding-worker-processes 1 until the pool adopts inherited listeners." + ) + get_engine_zmq_addresses = resolved + return get_engine_zmq_addresses, launch_core_engines + + logger = logging.getLogger(__name__) _ROLE_ENV = "DYN_VLLM_EMBEDDING_PROCESS_ROLE" @@ -463,7 +498,8 @@ def create_shared_embedding_engine_client( vllm_config.parallel_config._api_process_count = process_count vllm_config.parallel_config._api_process_rank = -1 executor_class = Executor.get_class(vllm_config) - addresses = get_engine_zmq_addresses(vllm_config, process_count) + engine_zmq_addresses, core_engines_launcher = _resolve_engine_utils() + addresses = engine_zmq_addresses(vllm_config, process_count) parent_logger_config = _client_config( vllm_config, @@ -476,7 +512,7 @@ def create_shared_embedding_engine_client( with ThreadPoolExecutor( max_workers=1, thread_name_prefix="dynamo-embedding-client" ) as client_pool: - with launch_core_engines( + with core_engines_launcher( vllm_config, executor_class, not disable_log_stats, diff --git a/components/src/dynamo/vllm/realtime/factories.py b/components/src/dynamo/vllm/realtime/factories.py index caedaf2ccc67..cb22a7cd80b4 100644 --- a/components/src/dynamo/vllm/realtime/factories.py +++ b/components/src/dynamo/vllm/realtime/factories.py @@ -89,7 +89,14 @@ def build_realtime_text_factories( consumed without starting speculative generation. """ from vllm.engine.protocol import StreamingInput - from vllm.entrypoints.chat_utils import load_chat_template + + try: + from vllm.entrypoints.chat_utils import load_chat_template + except ImportError: + # Newer vLLM moved chat_utils under vllm.renderers (renderers refactor, + # upstream main after 2026-10). Remove this fallback once every + # supported vLLM pin ships vllm.renderers.chat_utils. + from vllm.renderers.chat_utils.chat_template import load_chat_template from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat from vllm.entrypoints.serve.engine.protocol import ErrorResponse diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/frontend/protocol-field-handling.md b/docs/fern/pages/developer-guide/knowledge-base/modular-components/frontend/protocol-field-handling.md index c187324c1993..8998ccfb5ebb 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/frontend/protocol-field-handling.md +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/frontend/protocol-field-handling.md @@ -71,6 +71,7 @@ endpoint: | --- | --- | --- | --- | | `continue_final_message` | vLLM OpenAI-compatible extension | Dynamo-handled | Dynamo uses the value while preparing the chat template. The original field does not need to reach the engine after preprocessing. | | `temperature` | OpenAI Chat Completions API | Backend-handled | Backend sampling behavior can use a typed frontend representation. Typed storage does not make the sampling semantics frontend-owned. | +| `echo` | OpenAI Completions API | Jointly handled | The backend scores the prompt positions (its prompt logprobs); Dynamo rebuilds the echoed prompt text and per-token `logprobs` entries ahead of the completion from the tokens it sent. Unary responses prepend the prompt; a stream carries it in the chunk that delivers the engine's prompt logprobs. | | `prompt_logprobs` | vLLM OpenAI-compatible extension | Jointly handled | The backend computes the values while Dynamo carries, aggregates, and exposes the response payload. Document the streaming and non-streaming projections separately. | | `bad_words_token_ids` | Dynamo token-level extension | Backend-handled | On supported vLLM paths, Dynamo validates and transports the token sequences while vLLM applies the sampling constraint. | | `nvext.extra_fields` | Dynamo `nvext` extension | Dynamo-handled | This Dynamo-owned option selects response metadata. It is not a catch-all request map. | diff --git a/docs/fern/pages/use-cases/reinforcement-learning/integration-reference.md b/docs/fern/pages/use-cases/reinforcement-learning/integration-reference.md index faa87abd60a0..e2a7631ffc7c 100644 --- a/docs/fern/pages/use-cases/reinforcement-learning/integration-reference.md +++ b/docs/fern/pages/use-cases/reinforcement-learning/integration-reference.md @@ -251,7 +251,7 @@ A worker that uses a ModelExpress build with the RL cold-start policy declares t | [NeMo RL](nemo-rl.md) | Experimental | NeMo RL-managed Dynamo/vLLM fleet on Slurm and Ray | Fixed non-colocated fleet with framework-owned NCCL refit; not an external Dynamo deployment. | | [Slime](slime.md) | Experimental fixed worker set | Dynamo sidecars manage stock SGLang engines for [streaming external rollouts](https://github.com/THUDM/slime/pull/2272) | The [fixed-worker-set example](https://github.com/ai-dynamo/dynamo/blob/main/examples/README.md#integration-examples) sends generation through the Dynamo frontend. Slime uses static engine Services for control and weight updates. The example does not support elastic discovery or external-engine recovery. | | [Prime-RL](https://github.com/PrimeIntellect-ai/prime-rl) | Router available; full integration in progress | Prime-RL documents the Dynamo router as a drop-in option; a proposed Dynamo/vLLM sidecar adds worker discovery and external updates | Routing can be evaluated today, but the full adapter remains in upstream development and is not a released compatibility surface. | -| [OpenRLHF](https://github.com/OpenRLHF/OpenRLHF), [Miles](https://github.com/fleet-ai/miles-fleet), [SkyRL](https://github.com/NovaSky-AI/SkyRL), and [Polar](https://github.com/NVIDIA-NeMo/ProRL-Agent-Server) | No Dynamo guide | No maintained Dynamo adapter was found in the reviewed public sources | Add a guide only after a maintained integration completes the same generation, update, failure, and ownership checks. | +| [OpenRLHF](https://github.com/OpenRLHF/OpenRLHF), [Miles](https://github.com/radixark/miles), [SkyRL](https://github.com/NovaSky-AI/SkyRL), and [Polar](https://github.com/NVIDIA-NeMo/ProRL-Agent-Server) | No Dynamo guide | No maintained Dynamo adapter was found in the reviewed public sources | Add a guide only after a maintained integration completes the same generation, update, failure, and ownership checks. | A status reflects the documented integration, not whether the framework can call a generic HTTP endpoint. diff --git a/docs/fern/pages/use-cases/tool-calling-and-reasoning/reasoning-parsing.md b/docs/fern/pages/use-cases/tool-calling-and-reasoning/reasoning-parsing.md index 32808b61b723..537ba1154330 100644 --- a/docs/fern/pages/use-cases/tool-calling-and-reasoning/reasoning-parsing.md +++ b/docs/fern/pages/use-cases/tool-calling-and-reasoning/reasoning-parsing.md @@ -40,6 +40,7 @@ Some models emit reasoning separately from their final response. Dynamo can spli - `moonshotai/Kimi-K2.5*` / Kimi K2.6 format-compatible outputs: `--dyn-tool-call-parser kimi_k2 --dyn-reasoning-parser kimi_k25` - `google/gemma-4-*` thinking models: `--dyn-tool-call-parser gemma4 --dyn-reasoning-parser gemma4 --custom-jinja-template examples/chat_templates/gemma4_tool.jinja` - `Qwen/Qwen3.5*`: `--dyn-tool-call-parser qwen3_coder --dyn-reasoning-parser qwen3` + - Beam channel-grammar models: `--dyn-tool-call-parser beam --dyn-reasoning-parser beam` - MiniMax M2 style outputs: `--dyn-tool-call-parser minimax_m2 --dyn-reasoning-parser minimax_m2` - MiniMax M3 style outputs: `--dyn-tool-call-parser minimax_m3 --dyn-reasoning-parser minimax_m3` @@ -339,6 +340,20 @@ Parsers marked **Force-reasoning: Yes** emit reasoning content from token one wi + + + + **Parser:** `beam` + + **Upstream name:** none (Dynamo-native unified family) + + **Force-reasoning:** No + + **Notes:** `<\|channel\|>analysis<\|message\|>...<\|end\|>` is reasoning; the final channel is the answer. Pair with `--dyn-tool-call-parser beam` (both flags required). Supplied prior reasoning is always rendered back into the history (`preserve_thinking: true` and omitted are the same; `false` is accepted and has no effect); history turns without reasoning render an empty analysis block (`allow_missing_thinking` defaults to true on this route). + + + + diff --git a/docs/fern/pages/use-cases/tool-calling-and-reasoning/tool-call-parsing.mdx b/docs/fern/pages/use-cases/tool-calling-and-reasoning/tool-call-parsing.mdx index a0a61da9ff49..af9a19c4e49d 100644 --- a/docs/fern/pages/use-cases/tool-calling-and-reasoning/tool-call-parsing.mdx +++ b/docs/fern/pages/use-cases/tool-calling-and-reasoning/tool-call-parsing.mdx @@ -269,6 +269,20 @@ The **Upstream name** column shows where the vLLM or SGLang parser name differs + + + + **Parser:** `beam` + + **Upstream name:** none (Dynamo-native, served by the `dynamo-parsers-v2` unified family) + + **Notes:** Channel grammar: reasoning in `<\|channel\|>analysis<\|message\|>...<\|end\|>`, the answer in `<\|channel\|>final<\|message\|>...<\|return\|>`, and one `<\|start\|>assistant to=NAME<\|channel\|>analysis <\|constrain\|>xml<\|message\|>...<\|call\|>` message per tool call with `VALUE` arguments typed from the tool schema. Pair with `--dyn-reasoning-parser beam`; both flags are required and the pair routes to the unified parser without `DYN_PARSER_VERSION=2`. The route accepts the request field `preserve_thinking` (supplied prior assistant reasoning always stays in the rendered history; `true` and omitted are the same, `false` is accepted and has no effect — omit `reasoning_content` to drop it) and defaults `allow_missing_thinking=true` so history turns without reasoning render an empty analysis block instead of failing and `chat_template_kwargs.builtin_tools` (names resolved against `tools` for rendering only). + + **Per-turn grammar:** on the Beam route the frontend holds every chat turn to this channel grammar with a structural tag (a non-empty analysis channel, then either tool calls whose `rfl:parameter` values are held to their JSON-schema types or the final answer), mirroring the model builder's SGLang serving layer; `required` and named tool choices reuse the same tag, so they are enforced on the Harmony wire. A `response_format` (`json_object` or `json_schema`) is held inside the final channel behind the analysis, so the response keeps its reasoning and the channel structure the parser expects; a forced tool choice ignores `response_format`, and the `nvext` guided-decoding fields (`guided_json`, `guided_regex`, ...) stay bare constraints from the first token. Set `DYN_BEAM_TURN_GRAMMAR=0` on the frontend to disable it for automatic tool choice (forced choices keep theirs; `response_format` then falls back to the bare JSON constraint). A tool declared `strict: true` whose schema is outside OpenAI's strict subset is served as non-strict on this route (logged as a warning) instead of the HTTP 400 other routes return; the grammar holds every typeable tool either way. + + + + diff --git a/lib/bindings/kvbm/Cargo.lock b/lib/bindings/kvbm/Cargo.lock index 12aa44b9a375..a3484059b06e 100644 --- a/lib/bindings/kvbm/Cargo.lock +++ b/lib/bindings/kvbm/Cargo.lock @@ -1374,7 +1374,7 @@ version = "3.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "faf9468729b8cbcea668e36183cb69d317348c2e08e994829fb56ebfdfbaac34" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.48.0", ] [[package]] @@ -2320,9 +2320,8 @@ dependencies = [ [[package]] name = "dynamo-parsers-v2" -version = "0.7.21" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f43e47b6a68353cc04d4dfe19c69414a5884c9963b72c68937d019998c5ef72e" +version = "0.10.0" +source = "git+https://github.com/ai-dynamo/frontend-crates.git?rev=ef8f401393a730cae8ac5c96fe1d564dcf48a017#ef8f401393a730cae8ac5c96fe1d564dcf48a017" dependencies = [ "anyhow", "openai-harmony", @@ -2355,9 +2354,8 @@ dependencies = [ [[package]] name = "dynamo-renderer" -version = "8.0.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9dc3c1f04389ad78e3f9dad14360a191c27f3e3f574aba108b9986f5592cfea3" +version = "8.1.0" +source = "git+https://github.com/ai-dynamo/frontend-crates.git?rev=ef8f401393a730cae8ac5c96fe1d564dcf48a017#ef8f401393a730cae8ac5c96fe1d564dcf48a017" dependencies = [ "anyhow", "chrono", @@ -6083,7 +6081,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "be769465445e8c1474e9c5dac2018218498557af32d9ed057325ec9a41ae81bf" dependencies = [ "heck", - "itertools 0.14.0", + "itertools 0.13.0", "log", "multimap", "once_cell", @@ -6103,7 +6101,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "343d3bd7056eda839b03204e68deff7d1b13aba7af2b2fd16890697274262ee7" dependencies = [ "heck", - "itertools 0.14.0", + "itertools 0.13.0", "log", "multimap", "petgraph 0.8.3", @@ -6124,7 +6122,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a56d757972c98b346a9b766e3f02746cde6dd1cd1d1d563472929fdd74bec4d" dependencies = [ "anyhow", - "itertools 0.14.0", + "itertools 0.13.0", "proc-macro2", "quote", "syn 2.0.117", @@ -6137,7 +6135,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "27c6023962132f4b30eb4c172c91ce92d933da334c59c23cddee82358ddafb0b" dependencies = [ "anyhow", - "itertools 0.14.0", + "itertools 0.13.0", "proc-macro2", "quote", "syn 2.0.117", @@ -9094,7 +9092,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.48.0", ] [[package]] diff --git a/lib/bindings/python/Cargo.lock b/lib/bindings/python/Cargo.lock index 714482155aa4..a4d86af8d96a 100644 --- a/lib/bindings/python/Cargo.lock +++ b/lib/bindings/python/Cargo.lock @@ -1399,7 +1399,7 @@ version = "3.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "faf9468729b8cbcea668e36183cb69d317348c2e08e994829fb56ebfdfbaac34" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.48.0", ] [[package]] @@ -2372,9 +2372,8 @@ dependencies = [ [[package]] name = "dynamo-parsers-v2" -version = "0.7.21" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f43e47b6a68353cc04d4dfe19c69414a5884c9963b72c68937d019998c5ef72e" +version = "0.10.0" +source = "git+https://github.com/ai-dynamo/frontend-crates.git?rev=ef8f401393a730cae8ac5c96fe1d564dcf48a017#ef8f401393a730cae8ac5c96fe1d564dcf48a017" dependencies = [ "anyhow", "openai-harmony", @@ -2453,9 +2452,8 @@ dependencies = [ [[package]] name = "dynamo-renderer" -version = "8.0.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9dc3c1f04389ad78e3f9dad14360a191c27f3e3f574aba108b9986f5592cfea3" +version = "8.1.0" +source = "git+https://github.com/ai-dynamo/frontend-crates.git?rev=ef8f401393a730cae8ac5c96fe1d564dcf48a017#ef8f401393a730cae8ac5c96fe1d564dcf48a017" dependencies = [ "anyhow", "chrono", @@ -6523,7 +6521,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "be769465445e8c1474e9c5dac2018218498557af32d9ed057325ec9a41ae81bf" dependencies = [ "heck", - "itertools 0.14.0", + "itertools 0.13.0", "log", "multimap", "once_cell", @@ -6543,7 +6541,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "03da047801ff44bb6a4d407d4860c05fd70bb81714e6b2f3812603d5b145b042" dependencies = [ "heck", - "itertools 0.14.0", + "itertools 0.13.0", "log", "multimap", "petgraph 0.8.3", @@ -6564,7 +6562,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a56d757972c98b346a9b766e3f02746cde6dd1cd1d1d563472929fdd74bec4d" dependencies = [ "anyhow", - "itertools 0.14.0", + "itertools 0.13.0", "proc-macro2", "quote", "syn 2.0.118", @@ -6577,7 +6575,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b570b25f7617e43d59005d0990ccb79e950a423952cea19671b7a876da390adf" dependencies = [ "anyhow", - "itertools 0.14.0", + "itertools 0.13.0", "proc-macro2", "quote", "syn 2.0.118", @@ -9938,7 +9936,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.48.0", ] [[package]] diff --git a/lib/llm/benches/tokenizer_simple.rs b/lib/llm/benches/tokenizer_simple.rs index c2a9b54e84c2..07cf604bb218 100644 --- a/lib/llm/benches/tokenizer_simple.rs +++ b/lib/llm/benches/tokenizer_simple.rs @@ -62,7 +62,15 @@ pub fn decode(c: &mut Criterion) { let tokenizer: Arc = Arc::new(HuggingFaceTokenizer::from_file(TEST_TOKENIZER).unwrap()); let ds = DecodeStream::new(tokenizer, &[], false); - Decoder::new(ds, StopConditions::default(), false, false, None, None) + Decoder::new( + ds, + StopConditions::default(), + false, + false, + false, + None, + None, + ) }, |mut decoder| { for tok in black_box(TEST_TOKS) { @@ -86,7 +94,15 @@ pub fn decode_big(c: &mut Criterion) { let tokenizer: Arc = Arc::new(HuggingFaceTokenizer::from_file(TEST_TOKENIZER).unwrap()); let ds = DecodeStream::new(tokenizer, &[], false); - Decoder::new(ds, StopConditions::default(), false, false, None, None) + Decoder::new( + ds, + StopConditions::default(), + false, + false, + false, + None, + None, + ) }, |mut decoder| { for tok in black_box(&BIG_TEST_TOKS) { @@ -127,7 +143,15 @@ pub fn tiktoken_decode(c: &mut Criterion) { let tokenizer: Arc = Arc::new(TikTokenTokenizer::from_file_auto(TEST_TIKTOKEN).unwrap()); let ds = DecodeStream::new(tokenizer, &[], false); - Decoder::new(ds, StopConditions::default(), false, false, None, None) + Decoder::new( + ds, + StopConditions::default(), + false, + false, + false, + None, + None, + ) }, |mut decoder| { for tok in black_box(&toks) { diff --git a/lib/llm/src/backend.rs b/lib/llm/src/backend.rs index 555619841281..decfc7c85d38 100644 --- a/lib/llm/src/backend.rs +++ b/lib/llm/src/backend.rs @@ -72,6 +72,7 @@ struct DecoderUnfoldState { finished_choices: HashSet, validate_engine_decode: bool, no_stop_trim: bool, + retain_model_stop_token_text: bool, /// Set to true when all expected choices are finished locally, causing the stream to end finished: bool, /// Tokenizer used to decode top-logprob token_ids when the backend omits text @@ -136,6 +137,7 @@ struct DecoderParams { skip_special_tokens: bool, include_stop_str_in_output: bool, no_stop_trim: bool, + retain_model_stop_token_text: bool, tracker: Option>, n: u32, // In-process decoder checkpoint carried from a migrated attempt's last known-good @@ -160,6 +162,10 @@ impl DecoderParams { // mid-output) leak the token's text into `content` / `reasoning_content`. skip_special_tokens: request.output_options.skip_special_tokens.unwrap_or(true), no_stop_trim: request.output_options.no_stop_trim.unwrap_or(false), + retain_model_stop_token_text: request + .output_options + .retain_model_stop_token_text + .unwrap_or(false), include_stop_str_in_output: request .sampling_options .include_stop_str_in_output @@ -217,6 +223,7 @@ impl Backend { finished_choices: HashSet::new(), validate_engine_decode: self.validate_engine_decode, no_stop_trim: params.no_stop_trim, + retain_model_stop_token_text: params.retain_model_stop_token_text, finished: false, tokenizer: tokenizer.clone(), skip_special_tokens: params.skip_special_tokens, @@ -322,11 +329,12 @@ impl } // A worker may send already-trimmed text alongside the raw - // IDs. Re-decode those IDs for no_stop_trim instead of using - // the text fast path. Text-only output cannot honor this - // option; don't silently return trimmed or empty content. - // Empty stop trailers are valid after this choice emitted IDs. - if state.no_stop_trim + // IDs. Re-decode those IDs when stop-token text must survive + // (no_stop_trim, or a parser that needs the model's terminator) + // instead of using the text fast path. Text-only output cannot + // honor these options; don't silently return trimmed or empty + // content. Empty stop trailers are valid after this choice emitted IDs. + if (state.no_stop_trim || state.retain_model_stop_token_text) && let Some(data) = &output.data && data.token_ids.is_empty() && (data.text.as_ref().is_some_and(|text| !text.is_empty()) @@ -341,7 +349,7 @@ impl state.finished = true; let error = DynamoError::builder() .class(ErrorClass::BackendProtocol) - .message("no_stop_trim requires backend token IDs through the matched stop; the backend omitted them") + .message("no_stop_trim requires backend token IDs through the matched stop; the backend omitted them (also required when a parser retains the model's stop token text)") .build(); return Some((Annotated::from_err(error), state)); } @@ -350,7 +358,7 @@ impl if let Some(data) = &output.data && data.text.is_some() && !state.validate_engine_decode - && !state.no_stop_trim + && !(state.no_stop_trim || state.retain_model_stop_token_text) { // Text already decoded; track finish for this choice let choice_idx = data.index.unwrap_or(0); @@ -692,6 +700,10 @@ pub struct Decoder { // Include matched stop token text while preserving its stopping reason. no_stop_trim: bool, + // Include only the MODEL's hidden stop token text (system/EOS ids); user stop ids and + // user stop strings keep their normal trimming. A parser that needs the turn + // terminator sets this; `no_stop_trim` is the caller's broader opt-in. + retain_model_stop_token_text: bool, // text sequences that if found in the response will trigger a stop condition after the // minimum number of tokens have been generated (excluded from output) @@ -813,6 +825,7 @@ impl Decoder { stop_condition: StopConditions, include_stop_str_in_output: bool, no_stop_trim: bool, + retain_model_stop_token_text: bool, tracker: Option>, jail_seed: Option, ) -> Self { @@ -823,6 +836,7 @@ impl Decoder { stop_condition, include_stop_str_in_output, no_stop_trim, + retain_model_stop_token_text, tracker, jail_seed.unwrap_or_default(), Vec::new(), @@ -871,6 +885,7 @@ impl Decoder { params.stop_conditions.clone(), params.include_stop_str_in_output, params.no_stop_trim, + params.retain_model_stop_token_text, params.tracker.clone(), jail, pending_token_ids, @@ -886,6 +901,7 @@ impl Decoder { stop_condition: StopConditions, include_stop_str_in_output: bool, no_stop_trim: bool, + retain_model_stop_token_text: bool, tracker: Option>, jail: String, pending_token_ids: Vec, @@ -944,6 +960,7 @@ impl Decoder { visible_stop_ids, user_stop_ids, no_stop_trim, + retain_model_stop_token_text, hidden_stop_sequences, visible_stop_sequences, min_tokens: stop_condition.min_tokens.unwrap_or(0), @@ -967,11 +984,16 @@ impl Decoder { // increment the generated tokens self.generated_tokens += 1; + // The model's own stop token text stays visible when the route asked for it + // (`retain_model_stop_token_text`); user stop ids keep their normal trimming. + let retains_model_stop_text = + self.retain_model_stop_token_text && !self.user_stop_ids.contains(&token_id); // Finalize older text before dropping a hidden stop token. if !below_min_tokens && self.hidden_stop_ids.contains(&token_id) && !self.visible_stop_ids.contains(&token_id) && !self.no_stop_trim + && !retains_model_stop_text { let mut result = self.finish()?; result.token = None; @@ -1058,7 +1080,7 @@ impl Decoder { StopTrigger::HiddenStopTokenDetected(token_id) }; // Release any earlier withheld prefix before this stop token's own text. - let token = if self.no_stop_trim { + let token = if self.no_stop_trim || retains_model_stop_text { selected_token } else { None @@ -1991,6 +2013,7 @@ mod tests { stops, kind != "string", no_stop_trim, + false, None, None, ); @@ -2036,6 +2059,57 @@ mod tests { } } + #[test] + fn test_retain_model_stop_token_text_keeps_only_the_models_eos_text() { + let (stop_id, marker) = (104, "<|return|>"); + for kind in ["eos", "user", "string"] { + for streaming in [false, true] { + let tokenizer: Arc = Arc::new(CandidateDecoder); + let stops = StopConditions { + stop_token_ids: (kind == "user").then_some(vec![stop_id]), + stop_token_ids_hidden: (kind == "eos").then_some(vec![stop_id]), + stop: (kind == "string").then(|| vec![marker.to_string()]), + ..Default::default() + }; + let mut decoder = Decoder::new( + crate::tokenizers::DecodeStream::new(tokenizer, &[], false), + stops, + false, + false, + true, + None, + None, + ); + let mut text = if streaming { + decoder.process_token_ids(&[101]).unwrap().text.unwrap() + } else { + String::new() + }; + let ids = if streaming { + vec![stop_id, 999] + } else { + vec![101, stop_id, 999] + }; + let result = decoder.process_token_ids(&ids).unwrap(); + text.push_str(result.text.as_deref().unwrap_or_default()); + let expected = if kind == "eos" { + format!("Okay{marker}") + } else { + "Okay".to_string() + }; + assert_eq!(text, expected, "{kind} streaming={streaming}"); + match (kind, result.stop_trigger.unwrap()) { + ("eos", StopTrigger::HiddenStopTokenDetected(id)) + | ("user", StopTrigger::UserStopTokenDetected(id)) => assert_eq!(id, stop_id), + ("string", StopTrigger::HiddenStopSequenceDetected(seq)) => { + assert_eq!(seq, marker) + } + other => panic!("unexpected stop: {other:?}"), + } + } + } + } + /// When the tokenizer's decode() returns Err, Decoder::process_token_ids() /// should propagate the error. In the backend unfold closure, this error /// gets caught and converted to FinishReason::Error. @@ -2045,7 +2119,15 @@ mod tests { let decode_stream = crate::tokenizers::DecodeStream::new(tokenizer, &[], false); let stop_conditions = StopConditions::default(); - let mut decoder = Decoder::new(decode_stream, stop_conditions, false, false, None, None); + let mut decoder = Decoder::new( + decode_stream, + stop_conditions, + false, + false, + false, + None, + None, + ); let result = decoder.process_token_ids(&[42]); assert!( @@ -2068,7 +2150,15 @@ mod tests { let decode_stream = crate::tokenizers::DecodeStream::new(tokenizer, &[], false); let stop_conditions = StopConditions::default(); - let mut decoder = Decoder::new(decode_stream, stop_conditions, false, false, None, None); + let mut decoder = Decoder::new( + decode_stream, + stop_conditions, + false, + false, + false, + None, + None, + ); let result = decoder.process_token_ids(&[42]); let err = result.err().expect("should be Err"); diff --git a/lib/llm/src/discovery/watcher.rs b/lib/llm/src/discovery/watcher.rs index b2aa15bda31e..1bd21f101bc1 100644 --- a/lib/llm/src/discovery/watcher.rs +++ b/lib/llm/src/discovery/watcher.rs @@ -2758,6 +2758,7 @@ request_classifier: chat_template_args: None, thinking: None, thinking_token_budget: None, + preserve_thinking: None, media_io_kwargs: None, return_tokens_as_token_ids: None, unsupported_fields: Default::default(), diff --git a/lib/llm/src/entrypoint/input/text.rs b/lib/llm/src/entrypoint/input/text.rs index 8c078f9400cd..232ccd97830b 100644 --- a/lib/llm/src/entrypoint/input/text.rs +++ b/lib/llm/src/entrypoint/input/text.rs @@ -118,6 +118,7 @@ async fn main_loop( media_io_kwargs: None, return_tokens_as_token_ids: None, thinking_token_budget: None, + preserve_thinking: None, unsupported_fields: Default::default(), }; diff --git a/lib/llm/src/http/service.rs b/lib/llm/src/http/service.rs index a40cf0affecd..4f66b4ff06e8 100644 --- a/lib/llm/src/http/service.rs +++ b/lib/llm/src/http/service.rs @@ -215,10 +215,10 @@ mod tests { let result = apply_request_tool_call_parsing_options(options, &req).unwrap(); assert_eq!( result.reasoning_disabled, - if family == "gemma4" { - thinking != Some(true) - } else { - thinking == Some(false) + match family { + "beam" => false, + "gemma4" => thinking != Some(true), + _ => thinking == Some(false), }, "{selector} {tool:?} {reasoning:?} {thinking:?}" ); diff --git a/lib/llm/src/http/service/openai.rs b/lib/llm/src/http/service/openai.rs index fc06d021ef10..6048aaa3cbc0 100644 --- a/lib/llm/src/http/service/openai.rs +++ b/lib/llm/src/http/service/openai.rs @@ -1374,7 +1374,11 @@ async fn completions_single( .is_some_and(is_empty_completion_stream_response), ) }) - .map(move |response| { + .map(move |mut response| { + // `echo` rides the chunk as typed metadata up to this boundary. + if let Some(data) = response.data.as_mut() { + data.fold_echo(); + } let semantic_error = set_stream_semantic_error(&response, &producer_error_signal); // Calls observe_response() on each token let result = process_response_using_event_converter_and_observe_metrics( @@ -1654,6 +1658,10 @@ async fn completions_batch( for choice in &mut data.inner.choices { choice.index += prompt_idx_u32 * n_u32; } + // The echo names its choice by index too; keep them together. + if let Some(echo) = data.echo.as_mut() { + echo.choice_index += prompt_idx_u32 * n_u32; + } } response }); @@ -1724,7 +1732,11 @@ async fn completions_batch( .is_some_and(is_empty_completion_stream_response), ) }) - .map(move |response| { + .map(move |mut response| { + // `echo` rides the chunk as typed metadata up to this boundary. + if let Some(data) = response.data.as_mut() { + data.fold_echo(); + } let semantic_error = set_stream_semantic_error(&response, &producer_error_signal); // Calls observe_response() on each token let result = process_response_using_event_converter_and_observe_metrics( @@ -3334,7 +3346,7 @@ fn deduplicate_stream_roles( /// Completions variant of [`is_empty_stream_response`]. fn is_empty_completion_stream_response(resp: &NvCreateCompletionResponse) -> bool { - if resp.nvext.is_some() { + if resp.nvext.is_some() || resp.echo.is_some() { return false; } resp.inner.usage.is_none() @@ -3532,6 +3544,26 @@ async fn chat_completions( return Err(err_response); } + // Beam route: a `strict` tool whose schema is outside OpenAI's strict subset is + // served as non-strict instead of rejected (see + // `tool_choice::relax_noncompliant_strict_tools_for_beam`). The model lookup is + // repeated below with its own error handling, so a missing model is reported there. + if let Ok((_, model_parsing)) = state + .manager() + .get_chat_completions_engine_with_parsing(&model) + { + crate::preprocessor::tool_choice::declare_legacy_builtin_tools_for_beam( + &mut request, + model_parsing.tool_call_parser.as_deref(), + model_parsing.reasoning_parser.as_deref(), + ); + crate::preprocessor::tool_choice::relax_noncompliant_strict_tools_for_beam( + &mut request, + model_parsing.tool_call_parser.as_deref(), + model_parsing.reasoning_parser.as_deref(), + ); + } + // Handle Rest of Validation Errors if let Err(err_response) = validate_chat_completion_fields_generic(&request) { inflight_guard.mark_error(extract_error_type_from_response(&err_response)); @@ -8242,6 +8274,7 @@ mod tests { #[test] fn test_validate_chat_completion_required_fields_empty_messages() { let request = NvCreateChatCompletionRequest { + preserve_thinking: None, inner: CreateChatCompletionRequest { model: "test-model".to_string(), messages: vec![], @@ -8272,6 +8305,7 @@ mod tests { #[test] fn test_validate_chat_completion_required_fields_with_messages() { let request = NvCreateChatCompletionRequest { + preserve_thinking: None, inner: CreateChatCompletionRequest { model: "test-model".to_string(), messages: vec![ChatCompletionRequestMessage::User( @@ -8471,7 +8505,7 @@ mod tests { inner: CreateCompletionRequest { model: "test-model".to_string(), prompt: "Hello".into(), - logprobs: Some(6), + logprobs: Some(21), ..Default::default() }, common: Default::default(), @@ -8486,7 +8520,7 @@ mod tests { assert_eq!(error_response.0, StatusCode::BAD_REQUEST); assert_eq!( error_response.1.message, - format!("{VALIDATION_PREFIX}Logprobs must be between 0 and 5, got 6") + format!("{VALIDATION_PREFIX}Logprobs must be between 0 and 20, got 21") ); } } @@ -8526,6 +8560,7 @@ mod tests { fn test_bad_base_request_for_chatcompletion() { // Frequency Penalty: Should be a float between -2.0 and 2.0 let request = NvCreateChatCompletionRequest { + preserve_thinking: None, inner: CreateChatCompletionRequest { model: "test-model".to_string(), messages: vec![ChatCompletionRequestMessage::User( @@ -8559,6 +8594,7 @@ mod tests { // Presence Penalty: Should be a float between -2.0 and 2.0 let request = NvCreateChatCompletionRequest { + preserve_thinking: None, inner: CreateChatCompletionRequest { model: "test-model".to_string(), messages: vec![ChatCompletionRequestMessage::User( @@ -8591,6 +8627,7 @@ mod tests { // Temperature: Should be a float between 0.0 and 2.0 let request = NvCreateChatCompletionRequest { + preserve_thinking: None, inner: CreateChatCompletionRequest { model: "test-model".to_string(), messages: vec![ChatCompletionRequestMessage::User( @@ -8623,6 +8660,7 @@ mod tests { // Top P: Should be a float between 0.0 and 1.0 let request = NvCreateChatCompletionRequest { + preserve_thinking: None, inner: CreateChatCompletionRequest { model: "test-model".to_string(), messages: vec![ChatCompletionRequestMessage::User( @@ -8655,6 +8693,7 @@ mod tests { // Repetition Penalty: Should be a float between 0.0 and 2.0 let request = NvCreateChatCompletionRequest { + preserve_thinking: None, inner: CreateChatCompletionRequest { model: "test-model".to_string(), messages: vec![ChatCompletionRequestMessage::User( @@ -8689,6 +8728,7 @@ mod tests { // Top Logprobs: Should be a positive integer between 0 and 20 let request = NvCreateChatCompletionRequest { + preserve_thinking: None, inner: CreateChatCompletionRequest { model: "test-model".to_string(), messages: vec![ChatCompletionRequestMessage::User( @@ -10913,6 +10953,7 @@ mod tests { usage, }, nvext: None, + echo: None, } } @@ -10969,6 +11010,7 @@ mod tests { usage: Some(usage), }, nvext: None, + echo: None, } } diff --git a/lib/llm/src/preprocessor.rs b/lib/llm/src/preprocessor.rs index 59b5a8ce14cd..77c0d38b0448 100644 --- a/lib/llm/src/preprocessor.rs +++ b/lib/llm/src/preprocessor.rs @@ -198,6 +198,14 @@ pub struct PromptReasoningPrefill { legacy: bool, unified: bool, prefix: Option<&'static str>, + /// The prompt left the model's final (visible) channel open — an assistant + /// prefill with `continue_final_message` on a channel grammar (Beam) — so the + /// unified parser starts in `Response` instead of waiting for a channel opener. + response: bool, + /// Which channel the rendered prompt leaves open on the Beam route, which decides how + /// the per-turn grammar opens (`tool_choice`); `None` off the Beam route or without + /// prompt text, where the grammar assumes the bare header slot. + beam_channel: Option, } impl From for PromptReasoningPrefill { @@ -207,10 +215,33 @@ impl From for PromptReasoningPrefill { legacy: value, unified: value, prefix: None, + response: false, + beam_channel: None, } } } +/// Which Beam channel the rendered prompt leaves open after its last `<|start|>assistant` +/// header: `None` for the bare header slot (the model opens its own channel), `Reasoning` +/// for an open analysis message, `Response` for an open final message (assistant prefill). +fn beam_prompt_open_channel(prompt: &str) -> Option { + use dynamo_parsers_v2::UnifiedParserStartingState as State; + let tail = &prompt[prompt.rfind("<|start|>assistant")? + "<|start|>assistant".len()..]; + let open = |opener: &str, closers: &[&str]| { + tail.rfind(opener).is_some_and(|at| { + let rest = &tail[at + opener.len()..]; + !closers.iter().any(|closer| rest.contains(closer)) + }) + }; + if open("<|channel|>final<|message|>", &["<|return|>", "<|end|>"]) { + Some(State::Response) + } else if open("<|channel|>analysis<|message|>", &["<|end|>"]) { + Some(State::Reasoning) + } else { + Some(State::None) + } +} + impl PromptReasoningPrefill { fn for_route(self, route: &ToolProcessingRoute) -> bool { match route { @@ -1542,6 +1573,9 @@ struct NormalizedArgsRequest<'a, R> { inner: &'a R, normalize_tool_call_args: bool, continue_final_message: bool, + /// Beam's render-only view of `tools` / `chat_template_args` (`builtin_tools` + /// resolved against the declared tools); the request itself is never modified. + rendering: Option, } impl OAIChatLikeRequest for NormalizedArgsRequest<'_, R> { @@ -1574,7 +1608,10 @@ impl OAIChatLikeRequest for NormalizedArgsRequest<'_, R> } fn tools(&self) -> Option { - self.inner.tools() + match &self.rendering { + Some(rendering) => rendering.tools.clone(), + None => self.inner.tools(), + } } fn tool_choice(&self) -> Option { @@ -1594,7 +1631,10 @@ impl OAIChatLikeRequest for NormalizedArgsRequest<'_, R> } fn chat_template_args(&self) -> Option<&std::collections::HashMap> { - self.inner.chat_template_args() + match &self.rendering { + Some(rendering) => Some(&rendering.chat_template_args), + None => self.inner.chat_template_args(), + } } fn mm_processor_kwargs(&self) -> Option<&serde_json::Value> { @@ -1712,6 +1752,10 @@ pub struct OpenAIPreprocessor { tool_call_parser: Option, // Resolved at model setup so requests do not construct parsers to inspect their grammar. parser_requires_special_tokens: bool, + /// The model is served by the Beam unified family (`beam` tool-call and reasoning + /// parsers). Gates Beam-only request semantics: `preserve_thinking`, the + /// `builtin_tools` rendering view and stop-token visibility for the parser. + beam_route: bool, /// Normalize historical tool-call `function.arguments` from a JSON string /// to an object before MiniJinja rendering. Enabled for GLM-5.2 (glm47 /// parser) and any model that sets `normalize_tool_call_args: true` in its @@ -2170,6 +2214,75 @@ impl OpenAIPreprocessor { !matches!(reasoning_parser, Some("gpt_oss")) } + /// Beam-route history policy (the partner's serving contract). + /// + /// `preserve_thinking`: supplied prior assistant reasoning is always rendered back + /// into the prompt, also across tool calls; omitted and `true` are the same. + /// `false` is accepted for compatibility but has no effect: the template has no + /// switch for it, and a client that does not want prior reasoning in the prompt + /// simply omits `reasoning_content` from its history. Any other route would + /// silently ignore the field, so it is rejected there. + /// + /// Missing reasoning on a prior assistant turn is a quality guardrail in the Beam + /// template, not a parser requirement, so the Beam route defaults + /// `allow_missing_thinking=true` (the template then renders an empty analysis + /// block) unless the client set it explicitly. Reasoning that is supplied is never + /// dropped or filled in. + fn apply_preserve_thinking(&self, request: &mut NvCreateChatCompletionRequest) -> Result<()> { + if !self.beam_route { + if request.preserve_thinking.is_some() { + return Err(invalid_argument_error( + "`preserve_thinking` is only supported for models served with the `beam` parser", + )); + } + return Ok(()); + } + let args = request.chat_template_args.get_or_insert_with(HashMap::new); + args.entry("allow_missing_thinking".to_string()) + .or_insert(serde_json::Value::Bool(true)); + match request.preserve_thinking { + Some(true) => { + args.insert( + "preserve_thinking".to_string(), + serde_json::Value::Bool(true), + ); + } + Some(false) => { + tracing::debug!( + "preserve_thinking=false has no effect on the Beam route: supplied reasoning is kept; omit reasoning_content to drop it" + ); + } + None => {} + } + Ok(()) + } + + /// Families verified to have no supported non-thinking generation mode. + fn reasoning_is_mandatory(family: Option<&str>) -> bool { + matches!(family, Some("beam")) + } + + /// Validate the client's effective reasoning toggle against the model's policy. + fn validate_reasoning_request(&self, request: &NvCreateChatCompletionRequest) -> Result<()> { + let family = crate::protocols::openai::chat_completions::unified_parser::selected_family( + self.tool_call_parser.as_deref(), + self.runtime_config.reasoning_parser.as_deref(), + ); + if Self::reasoning_is_mandatory(family) + && Self::normalize_thinking_aliases(request.chat_template_args.as_ref(), family) + == Some(false) + { + let message = "The selected model does not support disabling reasoning. Remove the reasoning-off setting (thinking, enable_thinking, thinking_mode, or reasoning_effort=none)."; + return Err(DynamoError::builder() + .error_type(ErrorType::InvalidArgument) + .message(message) + .public_message(message) + .build() + .into()); + } + Ok(()) + } + pub(crate) fn has_structured_response_format(request: &R) -> bool { request.response_format().is_some_and(|format| { format @@ -2576,6 +2689,13 @@ impl OpenAIPreprocessor { let normalize_tool_call_args = mdc.runtime_config.tool_call_arguments_format == crate::local_model::runtime_config::ToolCallArgumentsFormat::JsonObject || mdc.runtime_config.tool_call_parser.as_deref() == Some("glm47"); + let beam_route = mdc.model_type.supports_chat() + && crate::protocols::openai::chat_completions::unified_parser::selected_family( + tool_call_parser.as_deref(), + mdc.runtime_config.reasoning_parser.as_deref(), + ) == Some( + crate::protocols::openai::chat_completions::unified_parser::BEAM_UNIFIED_FAMILY, + ); if let Some(ref lora_name) = lora_name { tracing::info!(model = %mdc.display_name, lora_name, "LoRA adapter detected in MDC"); @@ -2901,6 +3021,7 @@ impl OpenAIPreprocessor { kv_cache_block_size, tool_call_parser, parser_requires_special_tokens, + beam_route, normalize_tool_call_args, media_loader, token_budget, @@ -3110,7 +3231,8 @@ impl OpenAIPreprocessor { + SamplingOptionsProvider + StopConditionsProvider + OutputOptionsProvider - + NvExtProvider, + + NvExtProvider + + CommonExtProvider, >( &self, request: &R, @@ -3125,7 +3247,8 @@ impl OpenAIPreprocessor { + SamplingOptionsProvider + StopConditionsProvider + OutputOptionsProvider - + NvExtProvider, + + NvExtProvider + + CommonExtProvider, >( &self, request: &R, @@ -3199,6 +3322,15 @@ impl OpenAIPreprocessor { // for parsers that need special tokens preserved, unless the caller // has explicitly set `skip_special_tokens`. let mut output_options = request.extract_output_options()?; + // Beam's end-of-turn markers (`<|call|>`, `<|return|>`) are the model's stop + // tokens, and its parser completes a tool call only when it sees the marker. + // Keep the model's matched stop token text visible to the parser (it is consumed + // as control, never shown) while user stop strings and ids keep their normal + // trimming; not for a bare `response_format` document the engine's grammar owns + // in full (one the turn grammar does not hold inside the final channel). + if self.beam_route && !self.beam_output_is_bare_document(request) { + output_options.retain_model_stop_token_text = Some(true); + } if output_options.skip_special_tokens.is_none() && self.parser_requires_special_tokens { output_options.skip_special_tokens = Some(false); } else if Self::special_tokens_will_be_stripped( @@ -3420,15 +3552,27 @@ impl OpenAIPreprocessor { "Cannot continue the final message because the prompt formatter moves it before other messages", )); } - let formatted_prompt = if self.normalize_tool_call_args || continue_final { - self.apply_template_inner(&NormalizedArgsRequest { - inner: request, - normalize_tool_call_args: self.normalize_tool_call_args, - continue_final_message: continue_final, - })? + // Beam resolves `builtin_tools` names against `tools` for rendering only. + let rendering = if self.beam_route { + dynamo_renderer::beam::builtin_tools_view( + request, + self.runtime_config.exclude_tools_when_tool_choice_none, + ) + .map_err(|error| invalid_argument_error(error.to_string()))? } else { - self.apply_template_inner(request)? + None }; + let formatted_prompt = + if self.normalize_tool_call_args || continue_final || rendering.is_some() { + self.apply_template_inner(&NormalizedArgsRequest { + inner: request, + normalize_tool_call_args: self.normalize_tool_call_args, + continue_final_message: continue_final, + rendering, + })? + } else { + self.apply_template_inner(request)? + }; let Some(prompt) = formatted_prompt else { return Ok(None); }; @@ -5287,10 +5431,15 @@ impl OpenAIPreprocessor { tool_definitions, guided_tool_constraint, unified_parser::StreamPrefill { - starting_state: unified_parser::stream_prefill( - family, - prompt_injected_reasoning, - ), + starting_state: if prompt_prefill.response + && *family == unified_parser::BEAM_UNIFIED_FAMILY + { + // Assistant prefill inside Beam's final channel: the reply + // continues visible text and ends on `<|return|>`. + dynamo_parsers_v2::UnifiedParserStartingState::Response + } else { + unified_parser::stream_prefill(family, prompt_injected_reasoning) + }, prefix: prompt_prefill.prefix, }, family, @@ -6647,6 +6796,10 @@ impl OpenAIPreprocessor { prompt.ends_with("<|channel>thought") || prompt.ends_with("<|channel>") } Some("muse" | "muse_glimmer") => prompt.ends_with("assistant to=self<|message|>"), + Some("beam") => { + beam_prompt_open_channel(prompt) + == Some(dynamo_parsers_v2::UnifiedParserStartingState::Reasoning) + } _ => prompt.ends_with(""), } } @@ -6697,10 +6850,27 @@ impl OpenAIPreprocessor { } else { None }; + let beam_channel = + (crate::protocols::openai::chat_completions::unified_parser::configured_family( + tool_call_parser, + reasoning_parser, + ) == Some( + crate::protocols::openai::chat_completions::unified_parser::BEAM_UNIFIED_FAMILY, + )) + .then(|| { + formatted_prompt + .map(str::trim_end) + .and_then(beam_prompt_open_channel) + }) + .flatten(); + let response = + beam_channel == Some(dynamo_parsers_v2::UnifiedParserStartingState::Response); PromptReasoningPrefill { legacy, unified: unified || prefix.is_some(), prefix, + response, + beam_channel, } } @@ -6745,6 +6915,9 @@ impl OpenAIPreprocessor { default_thinking_mode: Option<&str>, unified_family: Option<&str>, ) -> bool { + if Self::reasoning_is_mandatory(unified_family) { + return false; + } let mut chat_template_args = request.chat_template_args.clone(); let from_client = Self::request_has_client_thinking_control(request); let constrained_generation = !from_client @@ -7617,10 +7790,12 @@ impl // Set stream=true for internal processing (after request payload capture) request.inner.stream = Some(true); + self.apply_preserve_thinking(&mut request)?; // Apply the deployment default before parser-specific normalization so // it can override an implicit model default (for example Kimi K2.5), // while explicit request controls still take precedence. let thinking_control_from_client = Self::request_has_client_thinking_control(&request); + self.validate_reasoning_request(&request)?; self.apply_default_thinking_mode(&mut request); Self::normalize_thinking_arg_with_source( &mut request, @@ -7629,6 +7804,16 @@ impl thinking_control_from_client, ); Self::normalize_kimi_k3_named_tool_choice(&mut request, self.tool_call_parser.as_deref()); + crate::preprocessor::tool_choice::declare_legacy_builtin_tools_for_beam( + &mut request, + self.tool_call_parser.as_deref(), + self.runtime_config.reasoning_parser.as_deref(), + ); + crate::preprocessor::tool_choice::relax_noncompliant_strict_tools_for_beam( + &mut request, + self.tool_call_parser.as_deref(), + self.runtime_config.reasoning_parser.as_deref(), + ); // create a response generator let response_generator = request.response_generator(context.id().to_string()); @@ -7659,7 +7844,7 @@ impl let guided_tool_constraint = self.apply_tool_choice_guided_decoding( &request, &mut common_request, - prompt_injected_reasoning.unified, + prompt_injected_reasoning, )?; let tool_processing_route = self.tool_processing_route(&request, &guided_tool_constraint)?; @@ -7904,6 +8089,25 @@ impl response_generator.update_isl(isl); } + // `echo: true` returns the prompt ahead of the completion. The generator + // rebuilds it from the tokens the engine saw once the engine reports its prompt + // logprobs (or on the choice's final chunk when it reports none). + if request.inner.echo == Some(true) && common_request.prompt_embeds.is_none() { + let prompt_text = match &request.inner.prompt { + dynamo_protocols::types::Prompt::String(text) => Some(text.clone()), + _ => None, + }; + response_generator.enable_echo( + Arc::clone(&common_request.token_ids), + prompt_text, + self.tokenizer.clone(), + common_request + .output_options + .skip_special_tokens + .unwrap_or(true), + ); + } + // repack the common completion request let common_request = context.map(|_| common_request); @@ -8415,7 +8619,7 @@ mod extra_args_media_copy_tests { use crate::model_card::ModelDeploymentCard; use crate::protocols::common::preprocessor::MultimodalData; - fn test_preprocessor() -> OpenAIPreprocessor { + pub(super) fn test_preprocessor() -> OpenAIPreprocessor { let mdc = ModelDeploymentCard::load_from_disk( "tests/data/sample-models/mock-llama-3.1-8b-instruct", None, @@ -9504,6 +9708,7 @@ mod tests { chat_template_args: None, thinking: None, thinking_token_budget: None, + preserve_thinking: None, media_io_kwargs: None, return_tokens_as_token_ids: None, unsupported_fields: Default::default(), @@ -13616,6 +13821,66 @@ mod tests { } } + #[tokio::test] + async fn beam_reasoning_off_is_rejected_before_backend_dispatch() { + let mut mdc = ModelDeploymentCard::load_from_disk( + "tests/data/sample-models/mock-llama-3.1-8b-instruct", + None, + ) + .unwrap(); + mdc.model_type = crate::model_type::ModelType::Chat; + mdc.runtime_config.tool_call_parser = Some("beam".to_string()); + mdc.runtime_config.reasoning_parser = Some("beam".to_string()); + let preprocessor = OpenAIPreprocessor::new(mdc).unwrap(); + + for extra in [ + serde_json::json!({"chat_template_kwargs": {"enable_thinking": false}}), + serde_json::json!({"chat_template_kwargs": {"thinking": false}}), + serde_json::json!({"chat_template_args": {"thinking": false}}), + serde_json::json!({"chat_template_kwargs": {"thinking_mode": "disabled"}}), + serde_json::json!({"thinking": {"type": "disabled"}}), + serde_json::json!({"reasoning_effort": "none"}), + serde_json::json!({"chat_template_kwargs": {"reasoning_effort": "none"}}), + ] { + for streaming in [false, true] { + let mut body = serde_json::json!({ + "model": "test-model", "messages": [{"role": "user", "content": "hi"}], + "stream": streaming + }); + body.as_object_mut() + .unwrap() + .extend(extra.as_object().unwrap().clone()); + let mut request: NvCreateChatCompletionRequest = + serde_json::from_value(body).unwrap(); + request.normalize_reasoning_template_args().unwrap(); + let error = match Operator::generate( + preprocessor.as_ref(), + PipelineContext::new(request), + Arc::new(UnreachableBackend), + ) + .await + { + Ok(_) => panic!("reasoning-off request must fail before dispatch"), + Err(error) => error, + }; + let error = error.downcast_ref::().unwrap(); + assert!(matches!( + crate::http::service::error::http_action_for_error(error), + crate::http::service::error::ClientErrorAction::Respond { + status: axum::http::StatusCode::BAD_REQUEST, + .. + } + )); + let message = error.public_message().unwrap(); + assert!( + message.contains("The selected model does not support disabling reasoning"), + "{message}" + ); + assert!(message.contains("Remove")); + } + } + } + #[tokio::test] async fn test_completion_operator_rejects_token_budget_overflow() { let mut mdc = ModelDeploymentCard::load_from_disk( @@ -13996,6 +14261,7 @@ mod tests { inner: &request, normalize_tool_call_args: true, continue_final_message: true, + rendering: None, }; assert!(request.typed_messages().is_some()); assert!(normalized.typed_messages().is_none()); @@ -14038,6 +14304,7 @@ mod tests { inner: request, normalize_tool_call_args: false, continue_final_message: true, + rendering: None, }) .unwrap() } else { @@ -15923,3 +16190,420 @@ mod tests { ])); } } + +#[cfg(test)] +mod beam_route_tests { + use super::*; + use crate::protocols::openai::chat_completions::NvCreateChatCompletionRequest; + use std::collections::HashMap; + + fn test_prompt_formatter_with_exclusion( + template: &str, + exclude_tools_when_tool_choice_none: bool, + ) -> Arc { + let template: dynamo_renderer::ChatTemplate = serde_json::from_value(serde_json::json!({ + "chat_template": template + })) + .unwrap(); + match dynamo_renderer::PromptFormatter::from_parts( + template, + dynamo_renderer::ContextMixins::default(), + exclude_tools_when_tool_choice_none, + ) + .unwrap() + { + dynamo_renderer::PromptFormatter::OAI(formatter) => formatter, + } + } + + fn synthetic_preprocessor(beam_route: bool, template: &str) -> OpenAIPreprocessor { + let mut preprocessor = extra_args_media_copy_tests::test_preprocessor(); + preprocessor.formatter = test_prompt_formatter_with_exclusion(template, true); + preprocessor + .runtime_config + .exclude_tools_when_tool_choice_none = true; + preprocessor.beam_route = beam_route; + preprocessor + } + + /// Beam's contract: `builtin_tools` names the tools from `tools` for the system + /// section; the rest go to the developer section. + const BUILTIN_NAMES_TEMPLATE: &str = "SYS:{% for t in tools %}{% if t.function.name in builtin_tools %}{{ t.function.name }}:{{ t.function.parameters.type }};{% endif %}{% endfor %}DEV:{% for t in tools %}{% if t.function.name not in builtin_tools %}{{ t.function.name }};{% endif %}{% endfor %}"; + + fn builtin_request( + tools: serde_json::Value, + builtin: serde_json::Value, + ) -> NvCreateChatCompletionRequest { + serde_json::from_value(serde_json::json!({ + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "tools": tools, + "chat_template_kwargs": {"builtin_tools": builtin}, + })) + .unwrap() + } + + fn function_tool(name: &str, schema_type: &str) -> serde_json::Value { + serde_json::json!({"type": "function", "function": { + "name": name, "description": "d", "parameters": {"type": schema_type} + }}) + } + + #[test] + fn beam_apply_template_declares_builtin_tools_by_name_from_the_whole_tool_list() { + let preprocessor = synthetic_preprocessor(true, BUILTIN_NAMES_TEMPLATE); + let request = builtin_request( + serde_json::json!([ + function_tool("bash", "string"), + function_tool("weather", "object") + ]), + serde_json::json!(["bash"]), + ); + let tools_before = serde_json::to_value(&request.inner.tools).unwrap(); + let rendered = preprocessor.apply_template(&request).unwrap().unwrap(); + assert_eq!(rendered.as_str(), "SYS:bash:string;DEV:weather;"); + assert_eq!( + serde_json::to_value(&request.inner.tools).unwrap(), + tools_before + ); + assert_eq!( + request.chat_template_args.as_ref().unwrap()["builtin_tools"], + serde_json::json!(["bash"]) + ); + } + + #[test] + fn beam_apply_template_renders_an_undeclared_builtin_object_as_a_tool() { + let preprocessor = synthetic_preprocessor(true, BUILTIN_NAMES_TEMPLATE); + let request = builtin_request( + serde_json::json!([ + function_tool("bash", "string"), + function_tool("weather", "object") + ]), + serde_json::json!([{"name": "python", "parameters": {"type": "string"}}, "bash"]), + ); + let rendered = preprocessor.apply_template(&request).unwrap().unwrap(); + assert_eq!( + rendered.as_str(), + "SYS:bash:string;python:string;DEV:weather;" + ); + } + + #[test] + fn beam_apply_template_rejects_undeclared_builtin_tool() { + let preprocessor = synthetic_preprocessor(true, BUILTIN_NAMES_TEMPLATE); + let request = builtin_request( + serde_json::json!([function_tool("weather", "object")]), + serde_json::json!(["bash"]), + ); + let err = preprocessor.apply_template(&request).unwrap_err(); + assert!(err.to_string().contains("not declared"), "{err}"); + } + + #[test] + fn non_beam_apply_template_leaves_builtin_tools_untouched() { + const TEMPLATE: &str = "{% for t in tools %}{{ t.function.name }};{% endfor %}"; + let preprocessor = synthetic_preprocessor(false, TEMPLATE); + let request = builtin_request( + serde_json::json!([ + function_tool("bash", "string"), + function_tool("weather", "object") + ]), + serde_json::json!(["bash"]), + ); + let rendered = preprocessor.apply_template(&request).unwrap().unwrap(); + assert_eq!(rendered.as_str(), "bash;weather;"); + } + + const HISTORY_TEMPLATE: &str = "{% for m in messages %}[{{ m.role }}:{{ m.content }}{% if m.reasoning_content %}|R={{ m.reasoning_content }}{% endif %}]{% endfor %}{% if preserve_thinking %}PT{% endif %}"; + + fn history_request(preserve_thinking: Option) -> NvCreateChatCompletionRequest { + let mut body = serde_json::json!({ + "model": "m", + "messages": [ + {"role": "user", "content": "q"}, + {"role": "assistant", "content": "a", "reasoning_content": "SENTINEL_THOUGHT"}, + {"role": "user", "content": "q2"}, + ], + }); + if let Some(value) = preserve_thinking { + body["preserve_thinking"] = serde_json::json!(value); + } + serde_json::from_value(body).unwrap() + } + + #[test] + fn beam_reasoning_validation_keeps_history_controls_and_other_families() { + let mut beam = synthetic_preprocessor(true, HISTORY_TEMPLATE); + beam.tool_call_parser = Some("beam".into()); + beam.runtime_config.reasoning_parser = Some("beam".into()); + for extra in [ + serde_json::json!({}), + serde_json::json!({"reasoning_effort": null}), + serde_json::json!({"chat_template_kwargs": {"reasoning_effort": null}}), + serde_json::json!({"reasoning_effort": "low"}), + serde_json::json!({"chat_template_kwargs": {"thinking": true}}), + serde_json::json!({"chat_template_kwargs": {"enable_thinking": true}}), + serde_json::json!({"chat_template_kwargs": {"thinking": true, "enable_thinking": false}}), + serde_json::json!({"chat_template_kwargs": {"preserve_thinking": false}}), + serde_json::json!({"chat_template_kwargs": {"allow_missing_thinking": false}}), + ] { + let mut body = serde_json::to_value(history_request(Some(false))).unwrap(); + body.as_object_mut() + .unwrap() + .extend(extra.as_object().unwrap().clone()); + let mut request: NvCreateChatCompletionRequest = serde_json::from_value(body).unwrap(); + request.normalize_reasoning_template_args().unwrap(); + beam.apply_preserve_thinking(&mut request).unwrap(); + beam.validate_reasoning_request(&request).unwrap(); + let rendered = beam.apply_template(&request).unwrap().unwrap(); + assert!(rendered.as_str().contains("SENTINEL_THOUGHT"), "{extra}"); + } + + let other = synthetic_preprocessor(false, HISTORY_TEMPLATE); + let mut request = history_request(None); + request.chat_template_args = Some(HashMap::from([( + "enable_thinking".into(), + serde_json::json!(false), + )])); + request.normalize_reasoning_template_args().unwrap(); + other.validate_reasoning_request(&request).unwrap(); + for (family, disabled) in [("beam", false), ("qwen3", true)] { + assert_eq!( + OpenAIPreprocessor::request_disables_reasoning_with_family( + &request, + Some(family), + Some(family), + None, + Some(family), + ), + disabled, + "{family}" + ); + } + } + + #[test] + fn preserve_thinking_beam_history_keeps_reasoning_for_none_and_true() { + let preprocessor = synthetic_preprocessor(true, HISTORY_TEMPLATE); + + let mut omitted = history_request(None); + preprocessor.apply_preserve_thinking(&mut omitted).unwrap(); + // The Beam route defaults the missing-reasoning guardrail off, nothing else. + let args = omitted.chat_template_args.as_ref().unwrap(); + assert_eq!(args["allow_missing_thinking"], serde_json::json!(true)); + assert!(!args.contains_key("preserve_thinking")); + let rendered = preprocessor.apply_template(&omitted).unwrap().unwrap(); + assert!(rendered.as_str().contains("|R=SENTINEL_THOUGHT")); + assert!(!rendered.as_str().ends_with("PT")); + + let mut requested = history_request(Some(true)); + requested.chat_template_args = Some(HashMap::from([( + "preserve_thinking".to_string(), + serde_json::json!(false), + )])); + preprocessor + .apply_preserve_thinking(&mut requested) + .unwrap(); + let rendered = preprocessor.apply_template(&requested).unwrap().unwrap(); + assert!(rendered.as_str().contains("|R=SENTINEL_THOUGHT")); + assert!(rendered.as_str().ends_with("PT")); + } + + #[test] + fn preserve_thinking_false_is_accepted_and_keeps_supplied_reasoning() { + let preprocessor = synthetic_preprocessor(true, HISTORY_TEMPLATE); + let mut request = history_request(Some(false)); + preprocessor.apply_preserve_thinking(&mut request).unwrap(); + let args = request.chat_template_args.as_ref().unwrap(); + assert!(!args.contains_key("preserve_thinking")); + assert_eq!(args["allow_missing_thinking"], serde_json::json!(true)); + let rendered = preprocessor.apply_template(&request).unwrap().unwrap(); + assert!( + rendered.as_str().contains("|R=SENTINEL_THOUGHT"), + "{rendered:?}" + ); + assert!(!rendered.as_str().ends_with("PT")); + } + + #[test] + fn beam_route_keeps_an_explicit_allow_missing_thinking() { + let preprocessor = synthetic_preprocessor(true, HISTORY_TEMPLATE); + let mut request = history_request(None); + request.chat_template_args = Some(HashMap::from([( + "allow_missing_thinking".to_string(), + serde_json::json!(false), + )])); + preprocessor.apply_preserve_thinking(&mut request).unwrap(); + assert_eq!( + request.chat_template_args.as_ref().unwrap()["allow_missing_thinking"], + serde_json::json!(false) + ); + } + + fn mock_card(parsers: Option<&str>) -> ModelDeploymentCard { + let mut card = ModelDeploymentCard::load_from_disk( + "tests/data/sample-models/mock-llama-3.1-8b-instruct", + None, + ) + .unwrap(); + card.model_type = crate::model_type::ModelType::Chat; + card.runtime_config.tool_call_parser = parsers.map(str::to_string); + card.runtime_config.reasoning_parser = parsers.map(str::to_string); + card + } + + fn weather_request() -> NvCreateChatCompletionRequest { + serde_json::from_value(serde_json::json!({ + "model": "test", + "messages": [{"role": "user", "content": "hi"}], + "tools": [{"type": "function", "function": { + "name": "get_weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}} + }}], + })) + .unwrap() + } + + /// Beam's `<|call|>` / `<|return|>` are the model's stop tokens and the parser needs + /// to see them, so the route keeps the matched stop token visible (and special + /// tokens decoded) unless the output is a bare `response_format` document. + #[test] + fn beam_route_keeps_stop_tokens_visible_unless_the_output_is_a_bare_document() { + let preprocessor = OpenAIPreprocessor::new(mock_card(Some("beam"))).unwrap(); + assert!(preprocessor.beam_route); + let request = weather_request(); + let output = preprocessor + .builder_with_lora(&request, None) + .unwrap() + .token_ids(vec![]) + .build() + .unwrap(); + assert_eq!( + output.output_options.retain_model_stop_token_text, + Some(true) + ); + assert_eq!(output.output_options.no_stop_trim, None); + assert_eq!(output.output_options.skip_special_tokens, Some(false)); + + // A `response_format` the turn grammar holds in the final channel keeps the + // channel structure, and a forced choice ignores `response_format`: both end on + // a marker the parser needs. + let json_object: dynamo_protocols::types::ResponseFormat = + serde_json::from_value(serde_json::json!({"type": "json_object"})).unwrap(); + for tool_choice in [None, Some(serde_json::json!("required"))] { + let mut structured = weather_request(); + structured.inner.response_format = Some(json_object.clone()); + structured.inner.tool_choice = tool_choice + .clone() + .map(|choice| serde_json::from_value(choice).unwrap()); + let output = preprocessor + .builder_with_lora(&structured, None) + .unwrap() + .token_ids(vec![]) + .build() + .unwrap(); + assert_eq!( + output.output_options.retain_model_stop_token_text, + Some(true), + "{tool_choice:?}" + ); + } + + // A bare document under an explicit guided-decoding constraint has no channel + // structure to keep. + let mut bare = weather_request(); + bare.inner.response_format = Some(json_object); + bare.common.guided_json = Some(serde_json::json!({"type": "object"})); + let output = preprocessor + .builder_with_lora(&bare, None) + .unwrap() + .token_ids(vec![]) + .build() + .unwrap(); + assert_eq!(output.output_options.retain_model_stop_token_text, None); + + let plain = OpenAIPreprocessor::new(mock_card(None)).unwrap(); + assert!(!plain.beam_route); + let output = plain + .builder_with_lora(&request, None) + .unwrap() + .token_ids(vec![]) + .build() + .unwrap(); + assert_eq!(output.output_options.retain_model_stop_token_text, None); + } + + #[test] + fn beam_prompt_tail_decides_the_parser_starting_state() { + use dynamo_parsers_v2::UnifiedParserStartingState as State; + const HISTORY: &str = "<|start|>user<|message|>q<|end|><|start|>assistant<|channel|>final<|message|>Hi<|return|><|start|>user<|message|>q2<|end|>"; + assert_eq!( + beam_prompt_open_channel(&format!("{HISTORY}<|start|>assistant")), + Some(State::None) + ); + assert_eq!( + beam_prompt_open_channel(&format!( + "{HISTORY}<|start|>assistant<|channel|>analysis<|message|>Let me" + )), + Some(State::Reasoning) + ); + assert_eq!( + beam_prompt_open_channel(&format!( + "{HISTORY}<|start|>assistant<|channel|>final<|message|>The planet is" + )), + Some(State::Response) + ); + assert_eq!( + beam_prompt_open_channel("<|start|>user<|message|>q<|end|>"), + None + ); + + let final_prefill = + format!("{HISTORY}<|start|>assistant<|channel|>final<|message|>The planet is"); + let prefill = OpenAIPreprocessor::prompt_reasoning_prefill_for_parsers( + Some("beam"), + Some("beam"), + Some(&final_prefill), + ); + assert!(prefill.response && !prefill.unified); + let thought_prefill = + format!("{HISTORY}<|start|>assistant<|channel|>analysis<|message|>Let me"); + let prefill = OpenAIPreprocessor::prompt_reasoning_prefill_for_parsers( + Some("beam"), + Some("beam"), + Some(&thought_prefill), + ); + assert!(prefill.unified && !prefill.response); + let plain = format!("{HISTORY}<|start|>assistant"); + let prefill = OpenAIPreprocessor::prompt_reasoning_prefill_for_parsers( + Some("beam"), + Some("beam"), + Some(&plain), + ); + assert!(!prefill.unified && !prefill.response); + // Another family's prompt never sets the Beam response prefill. + let prefill = OpenAIPreprocessor::prompt_reasoning_prefill_for_parsers( + Some("qwen3_coder"), + Some("qwen3"), + Some(&final_prefill), + ); + assert!(!prefill.response); + } + + #[test] + fn preserve_thinking_is_rejected_off_the_beam_route() { + let preprocessor = synthetic_preprocessor(false, HISTORY_TEMPLATE); + // Any supplied value is rejected; the value itself is not inspected. + let mut request = history_request(Some(true)); + let err = preprocessor + .apply_preserve_thinking(&mut request) + .unwrap_err(); + assert!(err.to_string().contains("preserve_thinking"), "{err}"); + + // Omitted is unchanged on a non-Beam route (no Beam defaults either). + let mut omitted = history_request(None); + preprocessor.apply_preserve_thinking(&mut omitted).unwrap(); + assert!(omitted.chat_template_args.is_none()); + } +} diff --git a/lib/llm/src/preprocessor/speculative_prefill.rs b/lib/llm/src/preprocessor/speculative_prefill.rs index e1c0322a5133..cb5538643cc8 100644 --- a/lib/llm/src/preprocessor/speculative_prefill.rs +++ b/lib/llm/src/preprocessor/speculative_prefill.rs @@ -611,6 +611,7 @@ mod tests { }); NvCreateChatCompletionRequest { + preserve_thinking: None, inner: CreateChatCompletionRequest { model: "mock-llama".to_string(), messages, diff --git a/lib/llm/src/preprocessor/structural_tag.rs b/lib/llm/src/preprocessor/structural_tag.rs index d3228acc7cb4..2c2a1caa9931 100644 --- a/lib/llm/src/preprocessor/structural_tag.rs +++ b/lib/llm/src/preprocessor/structural_tag.rs @@ -117,6 +117,19 @@ fn is_kimi_k2_parser(parser_name: Option<&str>) -> bool { parser_name == Some("kimi_k2") } +// Beam tool calls live on the Harmony wire: an analysis channel first, then the +// `<|start|>assistant to=NAME<|channel|>analysis <|constrain|>xml<|message|>` call +// header with an `rfl:parameter` body. The generic forced-tool JSON grammar replaces +// that whole wire with a bare JSON object, so the model emits no analysis channel and +// the response carries no `reasoning_content`; the Beam chat template then rejects the +// follow-up request whose history replays that tool-call turn without its reasoning. +// Forced choices therefore stay on the parser path; the preprocessor installs a regex +// over the native wire instead (see `tool_choice.rs`), which keeps the analysis channel +// and still forces the call header of an allowed tool. +pub(crate) fn is_beam_parser(parser_name: Option<&str>) -> bool { + parser_name == Some("beam") +} + fn requires_intrinsic_structural_tag(parser_name: Option<&str>, tool_choice: &ToolChoice) -> bool { // K2 forced calls and K3 named calls cannot use Dynamo's generic JSON-schema // fallback because both families emit native, marker-delimited formats. @@ -133,12 +146,14 @@ pub(super) fn requires_native_tool_call_format( tool_choice: &ToolChoice, ) -> bool { // The generic forced-tool JSON grammar describes a different wire format - // from Kimi's native marker-delimited calls. This also covers K3 required, - // which uses its prompt/parser path when structural guidance is off. - (is_kimi_k2_parser(parser_name) - && matches!(tool_choice, ToolChoice::Required | ToolChoice::Named(_))) - || (is_kimi_k3_parser(parser_name) - && matches!(tool_choice, ToolChoice::Required | ToolChoice::Named(_))) + // from Kimi's native marker-delimited calls and from Beam's Harmony calls. + // This also covers K3 required, which uses its prompt/parser path when + // structural guidance is off. + let forced = matches!(tool_choice, ToolChoice::Required | ToolChoice::Named(_)); + forced + && (is_kimi_k2_parser(parser_name) + || is_kimi_k3_parser(parser_name) + || is_beam_parser(parser_name)) } fn should_skip_tool_call_ban(exclude_tools_when_none: bool, tool_choice: &ToolChoice) -> bool { @@ -1239,6 +1254,24 @@ mod tests { )); } + #[test] + fn beam_forced_choices_keep_the_harmony_wire_and_its_analysis_channel() { + let named = ToolChoice::Named("get_weather".to_string()); + assert!(requires_native_tool_call_format( + Some("beam"), + &ToolChoice::Required + )); + assert!(requires_native_tool_call_format(Some("beam"), &named)); + assert!(!requires_native_tool_call_format( + Some("beam"), + &ToolChoice::Auto + )); + assert!(!requires_native_tool_call_format( + Some("beam"), + &ToolChoice::None + )); + } + #[test] fn kimi_forced_choices_keep_native_tags_when_global_mode_is_off() { let model_path = PathBuf::from(env!("CARGO_MANIFEST_DIR")) diff --git a/lib/llm/src/preprocessor/tool_choice.rs b/lib/llm/src/preprocessor/tool_choice.rs index a03f38c2c81a..455c0bb14f43 100644 --- a/lib/llm/src/preprocessor/tool_choice.rs +++ b/lib/llm/src/preprocessor/tool_choice.rs @@ -5,9 +5,10 @@ use std::borrow::Cow; -use crate::preprocessor::structural_tag::requires_native_tool_call_format; -use crate::preprocessor::{OpenAIPreprocessor, PreprocessedRequest}; +use crate::preprocessor::structural_tag::{is_beam_parser, requires_native_tool_call_format}; +use crate::preprocessor::{OpenAIPreprocessor, PreprocessedRequest, PromptReasoningPrefill}; use crate::protocols::openai::chat_completions::NvCreateChatCompletionRequest; +use crate::protocols::openai::common_ext::CommonExtProvider; use crate::protocols::openai::tools::{ ToolChoiceGuidance, get_tool_choice_guidance_from_tools, validate_openai_tool_choice, }; @@ -17,6 +18,7 @@ use dynamo_parsers::tool_calling::{ToolChoice, ToolDefinition}; use dynamo_protocols::types::{ ChatCompletionTool, ChatCompletionToolChoiceOption, CreateChatCompletionRequest, ResponseFormat, }; +use dynamo_renderer::OAIChatLikeRequest; use dynamo_runtime::error::{DynamoError, ErrorType}; fn invalid_argument(message: impl Into, public_message: Option<&str>) -> DynamoError { @@ -69,8 +71,9 @@ impl OpenAIPreprocessor { &self, request: &NvCreateChatCompletionRequest, common_request: &mut PreprocessedRequest, - prompt_injected_reasoning: bool, + prompt_prefill: impl Into, ) -> Result { + let prompt_prefill = prompt_prefill.into(); let tool_choice = request .inner .tool_choice @@ -97,12 +100,44 @@ impl OpenAIPreprocessor { )); } + // Beam: the parsers crate's turn grammar holds the whole turn to the response + // template over the Harmony wire — a non-empty analysis channel, then calls + // whose arguments keep their schema types (unheld, Beam writes integers as + // `N.0`) or the answer — so forced choices are enforced and the response keeps + // its reasoning. Dynamo only decides when it applies (`beam_turn_grammar_constraint`) + // and from which channel the rendered prompt leaves open. + let beam_constraint = if self.is_beam() { + beam_turn_grammar_constraint(request) + } else { + GuidedToolConstraint::None + }; + // Explicit guided decoding constrains assistant content and cannot be // composed with structural-tag tool calls. if !is_forced_tool_choice && has_explicit_guided_decoding { return Ok(GuidedToolConstraint::None); } + // Beam's `response_format` is held inside the final channel behind the + // analysis (beside the calls under automatic choice), the way the model + // builder's serving layer derives it, instead of a bare JSON document from + // the first token that loses the reasoning and the channel structure. + if !is_forced_tool_choice && beam_constraint.frames_structured_answer() { + if let Some(schema) = request.get_guided_json() + && let Some(structural_tag) = beam_turn_structural_tag( + tools.as_ref(), + &parser_tool_choice, + request.inner.parallel_tool_calls, + prompt_prefill.beam_channel, + Some(&schema), + ) + { + install_structural_tag(common_request, structural_tag); + return Ok(beam_constraint); + } + return Ok(GuidedToolConstraint::None); + } + if is_forced_tool_choice && has_response_format_constraint && let Some(gd) = common_request.sampling_options.guided_decoding.as_mut() @@ -115,13 +150,36 @@ impl OpenAIPreprocessor { &parser_tool_choice, &convert_tools(tools.as_ref()), request.inner.parallel_tool_calls, - prompt_injected_reasoning, + prompt_prefill.unified, request.inner.response_format.as_ref(), common_request, )? { return Ok(GuidedToolConstraint::StructuralTag); } + if matches!( + beam_constraint, + GuidedToolConstraint::BeamTurnGrammar { .. } + ) { + if let Some(structural_tag) = beam_turn_structural_tag( + tools.as_ref(), + &parser_tool_choice, + request.inner.parallel_tool_calls, + prompt_prefill.beam_channel, + None, + ) { + install_structural_tag(common_request, structural_tag); + return Ok(beam_constraint); + } + if is_forced_tool_choice { + tracing::debug!( + tool_choice = ?parser_tool_choice, + "Beam grammar could not hold this forced choice; using prompt/parser compatibility fallback" + ); + return Ok(GuidedToolConstraint::None); + } + } + // If the structural-tag layer did not compose this response format with // auto tool calls, it remains the sole assistant-output constraint. if !is_forced_tool_choice && has_response_format_constraint { @@ -203,15 +261,263 @@ impl OpenAIPreprocessor { } } -fn has_explicit_guided_decoding(request: &NvCreateChatCompletionRequest) -> bool { - request.common.guided_json.is_some() - || request.common.guided_regex.is_some() +impl OpenAIPreprocessor { + fn is_beam(&self) -> bool { + is_beam_parser(self.tool_call_parser.as_deref()) + || is_beam_parser(self.runtime_config.reasoning_parser.as_deref()) + } + + /// Whether the Beam turn grammar holds this request's `response_format` inside the + /// final channel behind the analysis (automatic or `none` tool choice, no explicit + /// guided decoding, grammar enabled) instead of a bare JSON document. + pub(crate) fn beam_frames_structured_response( + &self, + request: &R, + ) -> bool { + self.is_beam() && beam_turn_grammar_constraint(request).frames_structured_answer() + } + + /// Whether Beam's output is a bare `response_format` document the engine's grammar + /// owns in full, with no channel structure and so no end-of-turn marker for the + /// parser to see. A forced choice ignores `response_format` and generates native + /// calls; a framed `response_format` keeps the channel structure. + pub(crate) fn beam_output_is_bare_document( + &self, + request: &R, + ) -> bool { + Self::has_structured_response_format(request) + && !Self::has_guided_tool_choice(request) + && !self.beam_frames_structured_response(request) + } +} + +/// Beam route only: a tool declared `strict: true` whose schema lies outside +/// OpenAI's strict subset (every property required, `additionalProperties: false`, no +/// composition keywords — enforced upstream since #15153 with an HTTP 400) is served as +/// non-strict with a warning instead of rejected. The Beam turn grammar holds every +/// typeable tool regardless of `strict`, so nothing is lost, and clients that stamp +/// `strict: true` on every tool (inspect_ai's OpenAI-compatible provider does) keep +/// working. Other routes keep upstream's rejection. Returns the number of tools relaxed. +pub(crate) fn relax_noncompliant_strict_tools_for_beam( + request: &mut NvCreateChatCompletionRequest, + tool_call_parser: Option<&str>, + reasoning_parser: Option<&str>, +) -> usize { + if !(is_beam_parser(tool_call_parser) || is_beam_parser(reasoning_parser)) { + return 0; + } + let Some(tools) = request.inner.tools.as_mut() else { + return 0; + }; + let mut relaxed: Vec = Vec::new(); + let mut first_violation = None; + for tool in tools.iter_mut() { + if tool.function.strict != Some(true) { + continue; + } + if let Err(error) = + crate::protocols::openai::strict_schema::validate_strict_function(&tool.function) + { + // The name is unvalidated request data at this point; bound what is logged. + relaxed.push(tool.function.name.chars().take(64).collect()); + first_violation.get_or_insert_with(|| error.to_string()); + tool.function.strict = None; + } + } + if !relaxed.is_empty() { + // One line per request: a client that marks every tool strict would otherwise + // log once per tool per request. + tracing::warn!( + relaxed = relaxed.len(), + tools = %relaxed.iter().take(5).cloned().collect::>().join(", "), + first_violation = %first_violation.unwrap_or_default(), + "Beam route: tool(s) declared `strict` with a schema outside OpenAI's strict subset; serving them as non-strict (the turn grammar holds every typeable tool regardless)" + ); + } + relaxed.len() +} + +/// Beam route only: a `builtin_tools` entry given as a declaration object (the legacy +/// form; the 1.0 template takes names) that the request does not also declare in +/// `tools` is added to `tools`, so the prompt, the turn grammar and the parser all see +/// one tool set. Entries that are names, or objects already declared, are left to the +/// renderer, which validates them. Returns how many tools were added. +pub(crate) fn declare_legacy_builtin_tools_for_beam( + request: &mut NvCreateChatCompletionRequest, + tool_call_parser: Option<&str>, + reasoning_parser: Option<&str>, +) -> usize { + if !(is_beam_parser(tool_call_parser) || is_beam_parser(reasoning_parser)) { + return 0; + } + let Some(entries) = request + .chat_template_args + .as_ref() + .and_then(|args| args.get("builtin_tools")) + .and_then(serde_json::Value::as_array) + else { + return 0; + }; + let mut added: Vec = Vec::new(); + for entry in entries.iter().filter(|entry| entry.is_object()) { + let declaration = if entry.get("function").is_some() { + entry.clone() + } else { + serde_json::json!({"type": "function", "function": entry}) + }; + // A malformed declaration is left for the renderer's error message. + let Ok(tool) = serde_json::from_value::(declaration) else { + continue; + }; + let declared = request + .inner + .tools + .iter() + .flatten() + .chain(added.iter()) + .any(|existing| existing.function.name == tool.function.name); + if !declared { + added.push(tool); + } + } + let count = added.len(); + if count > 0 { + request + .inner + .tools + .get_or_insert_with(Vec::new) + .extend(added); + } + count +} + +fn install_structural_tag( + common_request: &mut PreprocessedRequest, + structural_tag: serde_json::Value, +) { + let gd = common_request + .sampling_options + .guided_decoding + .get_or_insert_default(); + gd.json = None; + gd.regex = None; + gd.structural_tag = Some(structural_tag); +} + +/// The constraint the Beam turn grammar implies for a request, decided from request +/// facts alone so a reconstructing frontend reports what preprocessing installed: +/// every turn is held unless `DYN_BEAM_TURN_GRAMMAR=0` opts automatic choice out or an +/// explicit guided-decoding field constrains the answer itself; a `response_format` on +/// an unforced choice is framed inside the final channel. (A forced choice with +/// explicit guided decoding is rejected at request time.) +/// Plain `ignore_eos` requests skip the turn grammar to avoid grammar-driven early stops. +fn beam_turn_grammar_constraint( + request: &R, +) -> GuidedToolConstraint { + let explicit = has_explicit_guided_decoding(request); + match ( + OpenAIPreprocessor::has_guided_tool_choice(request), + explicit, + ) { + (true, true) => GuidedToolConstraint::None, + (true, false) => GuidedToolConstraint::BeamTurnGrammar { + frames_answer: false, + }, + (false, _) if explicit || !beam_turn_grammar_enabled() => GuidedToolConstraint::None, + (false, _) + if request.common_ext().and_then(|common| common.ignore_eos) == Some(true) + && !OpenAIPreprocessor::has_structured_response_format(request) + && (request + .tool_choice() + .is_some_and(|choice| choice.as_str() == Some("none")) + || !OpenAIPreprocessor::request_has_effective_tools(request)) => + { + GuidedToolConstraint::None + } + (false, _) => GuidedToolConstraint::BeamTurnGrammar { + frames_answer: OpenAIPreprocessor::has_structured_response_format(request), + }, + } +} + +/// Environment variable that opts a deployment out of the per-turn Beam grammar for +/// automatic tool choice (`0`, `false`, `off` or `no`). Forced tool choices keep theirs: +/// without a grammar, `required` and named choices are not enforced on Beam's wire. +pub(crate) const BEAM_TURN_GRAMMAR_ENV: &str = "DYN_BEAM_TURN_GRAMMAR"; + +fn beam_turn_grammar_enabled() -> bool { + !dynamo_runtime::config::env_is_falsey(BEAM_TURN_GRAMMAR_ENV) +} + +/// The Beam turn grammar from the parsers crate's family registry, built for the channel +/// the rendered prompt leaves open (`None` = the bare header slot) and, with +/// `answer_schema`, holding the answer in the final channel to that JSON schema. `None` +/// where the grammar derives nothing (see +/// `dynamo_parsers_v2::structural_tag::builders::beam`). +fn beam_turn_structural_tag( + tools: &[ChatCompletionTool], + tool_choice: &ToolChoice, + parallel_tool_calls: Option, + prompt_channel: Option, + answer_schema: Option<&serde_json::Value>, +) -> Option { + use dynamo_parsers_v2::UnifiedParserStartingState as State; + use dynamo_parsers_v2::structural_tag::{ + PromptChannel, StructuralTagContext, StructuralTagSchemaMode, StructuralTagToolChoice, + }; + let builder = dynamo_parsers_v2::structural_tag_builder_for_family( + crate::protocols::openai::chat_completions::unified_parser::BEAM_UNIFIED_FAMILY, + )?; + let v2_tools: Vec = tools + .iter() + .map(|tool| dynamo_parsers_v2::Tool { + name: tool.function.name.clone(), + description: None, + parameters: tool + .function + .parameters + .clone() + .unwrap_or(serde_json::Value::Null), + strict: tool.function.strict, + }) + .collect(); + let context = StructuralTagContext { + tool_choice: match tool_choice { + ToolChoice::None => StructuralTagToolChoice::None, + ToolChoice::Auto => StructuralTagToolChoice::Auto, + ToolChoice::Required => StructuralTagToolChoice::Required, + ToolChoice::Named(name) => StructuralTagToolChoice::Named(name.as_str()), + }, + tools: &v2_tools, + parallel_tool_calls, + schema_mode: StructuralTagSchemaMode::Auto, + structured_output_schema: answer_schema, + starts_in_reasoning: false, + }; + let channel = match prompt_channel { + Some(State::Reasoning) => PromptChannel::OpenReasoning, + Some(State::Response) => PromptChannel::OpenAnswer, + Some(State::None) | None => PromptChannel::MessageStart, + }; + match builder.build_turn(&context, channel) { + Ok(tag) => tag, + Err(error) => { + tracing::warn!(%error, "Beam turn grammar failed to build; leaving the turn unconstrained"); + None + } + } +} + +/// An `nvext`/top-level guided-decoding field on the request itself; `get_guided_json` +/// is not consulted because it also derives a schema from `response_format`. +fn has_explicit_guided_decoding(request: &R) -> bool { + request + .common_ext() + .is_some_and(|common| common.guided_json.is_some() || common.guided_grammar.is_some()) + || request.get_guided_regex().is_some() || request - .common - .guided_choice - .as_ref() - .is_some_and(|v| !v.is_empty()) - || request.common.guided_grammar.is_some() + .get_guided_choice() + .is_some_and(|choices| !choices.is_empty()) } fn has_response_format_constraint(request: &NvCreateChatCompletionRequest) -> bool { @@ -310,6 +616,12 @@ pub(crate) fn guided_tool_constraint_with_effective_tools( .unwrap_or(&ChatCompletionToolChoiceOption::Auto); validate_openai_tool_choice(Some(tool_choice), Some(tools)) .map_err(|error| invalid_argument(error.to_string(), None))?; + // Beam holds every turn with its grammar (`apply_tool_choice_guided_decoding`); a + // reconstructing frontend reports the same so parsing frames the arguments and a + // `response_format` answer the way that grammar shaped them. + if is_beam_parser(tool_call_parser) || is_beam_parser(reasoning_parser) { + return Ok(beam_turn_grammar_constraint(request)); + } let is_forced_tool_choice = matches!( tool_choice, ChatCompletionToolChoiceOption::Required | ChatCompletionToolChoiceOption::Named(_) @@ -325,7 +637,7 @@ pub(crate) fn guided_tool_constraint_with_effective_tools( return Ok(GuidedToolConstraint::None); } let parser_tool_choice = convert_tool_choice(tool_choice); - // Kimi forced requests use marker-delimited native formats that generic JSON + // Kimi and Beam forced requests use native wire formats that generic JSON // guidance cannot represent. Keep this compatibility path aligned with // request preprocessing for either configured parser slot. if requires_native_tool_call_format(tool_call_parser, &parser_tool_choice) @@ -402,6 +714,10 @@ mod tests { }]) } + fn preprocessor_for(parser: &str) -> Arc { + preprocessor(parser, StructuralTagMode::Off) + } + fn preprocessor(parser: &str, mode: StructuralTagMode) -> Arc { let model_path = PathBuf::from(env!("CARGO_MANIFEST_DIR")) .join("tests/data/sample-models/mock-llama-3.1-8b-instruct"); @@ -760,6 +1076,540 @@ mod tests { } } + #[test] + fn beam_forced_choices_install_the_harmony_wire_structural_tag_and_no_json_schema() { + // Single tool, `required`, parallel calls allowed by default: the turn grammar the + // parsers crate derives for the same tools (asserted there against the model + // builder's own fixtures) is what the policy installs, and nothing else. + let first = request(json!({"tools": tools(), "tool_choice": "required"})); + let blueprint = beam_turn_structural_tag( + effective_tools(&first.inner).unwrap().as_ref(), + &ToolChoice::Required, + None, + None, + None, + ) + .expect("the parsers crate derives the Beam turn grammar"); + let preprocessor = preprocessor("beam", StructuralTagMode::Off); + let mut common_request = preprocessed_request(); + + let applied = preprocessor + .apply_tool_choice_guided_decoding(&first, &mut common_request, false) + .unwrap(); + + assert_eq!( + applied, + GuidedToolConstraint::BeamTurnGrammar { + frames_answer: false + } + ); + let guided = common_request + .sampling_options + .guided_decoding + .as_ref() + .expect("beam forced choice installs guided decoding"); + assert!(guided.json.is_none() && guided.regex.is_none()); + assert_eq!( + guided.structural_tag.as_ref().expect("structural tag"), + &blueprint + ); + + // A named choice among several tools forces that tool only; `required` offers each + // tool, and `parallel_tool_calls: false` drops the repeated-call group. + let mut tools = tools(); + tools.as_array_mut().unwrap().push(json!({ + "type": "function", + "function": {"name": "get-time", "parameters": {"type": "object", "properties": {}}} + })); + for (tool_choice, expected, unexpected) in [ + ( + json!({"type": "function", "function": {"name": "get-time"}}), + vec!["to=get-time<"], + vec!["to=get_weather<"], + ), + ( + json!("required"), + vec!["to=get_weather<", "to=get-time<"], + vec![], + ), + ] { + let forced = request(json!({ + "tools": tools, + "tool_choice": tool_choice, + "parallel_tool_calls": false + })); + let mut common_request = preprocessed_request(); + preprocessor + .apply_tool_choice_guided_decoding(&forced, &mut common_request, false) + .unwrap(); + let tag = common_request + .sampling_options + .guided_decoding + .as_ref() + .and_then(|guided| guided.structural_tag.as_ref()) + .expect("structural tag") + .to_string(); + for needle in expected { + assert!(tag.contains(needle), "{tag}"); + } + for needle in unexpected { + assert!(!tag.contains(needle), "{tag}"); + } + assert!( + !tag.contains( + r#""type":"star","content":{"type":"sequence","elements":[{"type":"or""# + ), + "{tag}" + ); + } + } + + fn installed_structural_tag(common_request: &PreprocessedRequest) -> Option<&Value> { + common_request + .sampling_options + .guided_decoding + .as_ref() + .and_then(|guided| guided.structural_tag.as_ref()) + } + + #[test] + fn beam_ignore_eos_skips_only_the_plain_chat_turn_grammar() { + for mode in [StructuralTagMode::Off, StructuralTagMode::On] { + let preprocessor = preprocessor("beam", mode); + for extra in [ + json!({}), + json!({"tools": tools(), "tool_choice": "none"}), + json!({"tool_choice": "auto", "response_format": {"type": "text"}}), + ] { + for ignore_eos in [None, Some(false), Some(true)] { + let mut request = request(extra.clone()); + request.common.ignore_eos = ignore_eos; + let mut common_request = preprocessed_request(); + let applied = preprocessor + .apply_tool_choice_guided_decoding(&request, &mut common_request, false) + .unwrap(); + let expected = if ignore_eos == Some(true) { + GuidedToolConstraint::None + } else { + GuidedToolConstraint::BeamTurnGrammar { + frames_answer: false, + } + }; + assert_eq!(applied, expected, "{mode:?} {extra} {ignore_eos:?}"); + assert_eq!( + installed_structural_tag(&common_request).is_some(), + ignore_eos != Some(true) + ); + // Reconstructed response policy must match preprocessing with either parser slot. + for parsers in [(Some("beam"), None), (None, Some("beam"))] { + assert_eq!( + guided_tool_constraint(&request, parsers.0, parsers.1, false).unwrap(), + expected, + ); + } + } + } + } + } + + #[test] + fn beam_ignore_eos_preserves_tool_and_output_constraints() { + let preprocessor = preprocessor_for("beam"); + for extra in [ + json!({"tools": tools()}), + json!({"messages": [ + {"role": "system", "content": "", "tools": [{"name": "lookup", "parameters": {"type": "object"}}]}, + {"role": "user", "content": "hi"} + ]}), + json!({"response_format": {"type": "json_object"}}), + json!({"response_format": {"type": "json_schema", "json_schema": {"name": "answer", "schema": {"type": "object"}}}}), + ] { + let mut request = request(extra.clone()); + let mut normal = preprocessed_request(); + let expected = preprocessor + .apply_tool_choice_guided_decoding(&request, &mut normal, false) + .unwrap(); + assert!(installed_structural_tag(&normal).is_some(), "{extra}"); + request.common.ignore_eos = Some(true); + let mut ignored = preprocessed_request(); + let applied = preprocessor + .apply_tool_choice_guided_decoding(&request, &mut ignored, false) + .unwrap(); + assert_eq!(applied, expected, "{extra}"); + assert_eq!( + installed_structural_tag(&ignored), + installed_structural_tag(&normal), + "{extra}" + ); + } + } + + #[test] + fn beam_automatic_choice_installs_the_turn_grammar_and_keeps_the_parser_route() { + let preprocessor = preprocessor("beam", StructuralTagMode::Off); + let auto = request(json!({"tools": tools()})); + let blueprint = beam_turn_structural_tag( + effective_tools(&auto.inner).unwrap().as_ref(), + &ToolChoice::Auto, + None, + None, + None, + ) + .expect("the parsers crate derives the Beam turn grammar"); + let mut common_request = preprocessed_request(); + + let applied = preprocessor + .apply_tool_choice_guided_decoding(&auto, &mut common_request, false) + .unwrap(); + + assert_eq!( + applied, + GuidedToolConstraint::BeamTurnGrammar { + frames_answer: false + } + ); + let guided = common_request + .sampling_options + .guided_decoding + .as_ref() + .expect("automatic choice installs the turn grammar"); + assert!(guided.json.is_none() && guided.regex.is_none()); + assert_eq!( + guided.structural_tag.as_ref().expect("structural tag"), + &blueprint + ); + + // Plain chat and `tool_choice: none` hold the reasoning, then the answer. + let answer_only = beam_turn_structural_tag(&[], &ToolChoice::Auto, None, None, None) + .expect("plain chat is held to analysis then the answer"); + for extra in [json!({}), json!({"tools": tools(), "tool_choice": "none"})] { + let mut common_request = preprocessed_request(); + preprocessor + .apply_tool_choice_guided_decoding(&request(extra), &mut common_request, false) + .unwrap(); + assert_eq!( + installed_structural_tag(&common_request), + Some(&answer_only) + ); + } + + // An explicit guided-decoding field constrains the answer instead of the turn. + let mut constrained = request(json!({"tools": tools()})); + constrained.common.guided_regex = Some("[a-z]+".to_string()); + let mut common_request = preprocessed_request(); + let applied = preprocessor + .apply_tool_choice_guided_decoding(&constrained, &mut common_request, false) + .unwrap(); + assert_eq!(applied, GuidedToolConstraint::None); + assert!(installed_structural_tag(&common_request).is_none()); + + // Other parsers see no Beam grammar. + let hermes = preprocessor_for("hermes"); + let mut common_request = preprocessed_request(); + hermes + .apply_tool_choice_guided_decoding( + &request(json!({"tools": tools()})), + &mut common_request, + false, + ) + .unwrap(); + assert!(installed_structural_tag(&common_request).is_none()); + } + + #[test] + fn beam_response_format_is_held_in_the_final_channel_behind_the_analysis() { + let preprocessor = preprocessor("beam", StructuralTagMode::Off); + let schema = json!({ + "type": "object", + "properties": {"ok": {"type": "boolean"}}, + "required": ["ok"] + }); + let format = json!({ + "type": "json_schema", + "json_schema": {"name": "answer", "schema": schema} + }); + + // Automatic choice: the calls stay available beside the framed answer, and the + // bare JSON constraint `response_format` arrived as is replaced. + let auto = request(json!({"tools": tools(), "response_format": format})); + assert!(preprocessor.beam_frames_structured_response(&auto)); + assert!(!preprocessor.beam_output_is_bare_document(&auto)); + let blueprint = beam_turn_structural_tag( + effective_tools(&auto.inner).unwrap().as_ref(), + &ToolChoice::Auto, + None, + None, + Some(&schema), + ) + .expect("the parsers crate derives the framed answer"); + let mut common_request = preprocessed_request(); + common_request + .sampling_options + .guided_decoding + .get_or_insert_default() + .json = Some(schema.clone()); + let applied = preprocessor + .apply_tool_choice_guided_decoding(&auto, &mut common_request, false) + .unwrap(); + assert_eq!( + applied, + GuidedToolConstraint::BeamTurnGrammar { + frames_answer: true + } + ); + let guided = common_request + .sampling_options + .guided_decoding + .as_ref() + .expect("the framed answer is a structural tag"); + assert!(guided.json.is_none() && guided.regex.is_none()); + assert_eq!(guided.structural_tag.as_ref().expect("tag"), &blueprint); + let tag = blueprint.to_string(); + assert!( + tag.contains("<|channel|>final<|message|>") + && tag.contains(r#""type":"json_schema""#) + && tag.contains(r#""ok""#) + && tag.contains("to=get_weather<"), + "{tag}" + ); + + // `json_object` is any object; `tool_choice: none` and no tools frame the + // answer only. + for extra in [ + json!({"response_format": {"type": "json_object"}}), + json!({ + "tools": tools(), + "tool_choice": "none", + "response_format": {"type": "json_object"} + }), + ] { + let answer_only = request(extra); + let mut common_request = preprocessed_request(); + let applied = preprocessor + .apply_tool_choice_guided_decoding(&answer_only, &mut common_request, false) + .unwrap(); + assert_eq!( + applied, + GuidedToolConstraint::BeamTurnGrammar { + frames_answer: true + } + ); + let tag = installed_structural_tag(&common_request) + .expect("framed answer") + .to_string(); + assert!( + tag.contains(r#""json_schema":{"type":"object"}"#) + && !tag.contains("to=get_weather<"), + "{tag}" + ); + } + + // A forced choice ignores `response_format` (it governs content, not calls) and + // still generates native calls that end on a marker. + let forced = request(json!({ + "tools": tools(), + "tool_choice": "required", + "response_format": {"type": "json_object"} + })); + assert!(!preprocessor.beam_frames_structured_response(&forced)); + assert!(!preprocessor.beam_output_is_bare_document(&forced)); + let mut common_request = preprocessed_request(); + let applied = preprocessor + .apply_tool_choice_guided_decoding(&forced, &mut common_request, false) + .unwrap(); + assert_eq!( + applied, + GuidedToolConstraint::BeamTurnGrammar { + frames_answer: false + } + ); + let tag = installed_structural_tag(&common_request) + .expect("forced calls") + .to_string(); + assert!( + !tag.contains("json_schema") && tag.contains("to=get_weather<"), + "{tag}" + ); + + // Explicit guided decoding stays a bare constraint from the first token. + let mut explicit = request(json!({ + "tools": tools(), + "response_format": {"type": "json_object"} + })); + explicit.common.guided_json = Some(schema.clone()); + assert!(!preprocessor.beam_frames_structured_response(&explicit)); + assert!(preprocessor.beam_output_is_bare_document(&explicit)); + let mut common_request = preprocessed_request(); + let applied = preprocessor + .apply_tool_choice_guided_decoding(&explicit, &mut common_request, false) + .unwrap(); + assert_eq!(applied, GuidedToolConstraint::None); + assert!(installed_structural_tag(&common_request).is_none()); + + // Other parsers leave `response_format` to the engine. + let hermes = preprocessor_for("hermes"); + assert!(!hermes.beam_frames_structured_response(&auto)); + let mut common_request = preprocessed_request(); + let applied = hermes + .apply_tool_choice_guided_decoding(&auto, &mut common_request, false) + .unwrap(); + assert_eq!(applied, GuidedToolConstraint::None); + assert!(installed_structural_tag(&common_request).is_none()); + + // A reconstructing frontend reports what preprocessing installed. + for (request, parsers, expected) in [ + ( + &auto, + (Some("beam"), None), + GuidedToolConstraint::BeamTurnGrammar { + frames_answer: true, + }, + ), + ( + &forced, + (Some("beam"), Some("beam")), + GuidedToolConstraint::BeamTurnGrammar { + frames_answer: false, + }, + ), + (&explicit, (None, Some("beam")), GuidedToolConstraint::None), + (&auto, (Some("hermes"), None), GuidedToolConstraint::None), + ] { + assert_eq!( + guided_tool_constraint(request, parsers.0, parsers.1, false).unwrap(), + expected, + "{parsers:?}" + ); + } + } + + #[test] + fn beam_route_declares_legacy_builtin_declarations_as_tools() { + let weather = json!({"type": "function", "function": { + "name": "weather", "parameters": {"type": "object", "properties": {}}}}); + let mut beam = request(json!({"tools": [weather.clone()]})); + beam.chat_template_args = Some(std::collections::HashMap::from([( + "builtin_tools".to_string(), + json!([ + "weather", + {"name": "bash", "parameters": {"type": "object", "properties": { + "cmd": {"type": "string"}}}}, + {"type": "function", "function": {"name": "weather"}}, + {"type": "function", "function": {"name": "python", + "parameters": {"type": "object", "properties": {}}}} + ]), + )])); + + // Names and already-declared objects add nothing; undeclared objects join `tools` + // in OpenAI form, once each. + assert_eq!( + declare_legacy_builtin_tools_for_beam(&mut beam, Some("beam"), None), + 2 + ); + let names: Vec<&str> = beam + .inner + .tools + .as_ref() + .unwrap() + .iter() + .map(|tool| tool.function.name.as_str()) + .collect(); + assert_eq!(names, vec!["weather", "bash", "python"]); + assert_eq!( + declare_legacy_builtin_tools_for_beam(&mut beam, Some("beam"), None), + 0 + ); + + // Other routes do not read `builtin_tools`. + let mut hermes = request(json!({"tools": [weather]})); + hermes.chat_template_args = beam.chat_template_args.clone(); + assert_eq!( + declare_legacy_builtin_tools_for_beam(&mut hermes, Some("hermes"), None), + 0 + ); + assert_eq!(hermes.inner.tools.as_ref().unwrap().len(), 1); + } + + #[test] + fn beam_route_serves_noncompliant_strict_tools_as_non_strict() { + use crate::engines::ValidateRequest; + // BFCL-style: optional properties, so `required` does not list every property. + let loose = json!({"type": "function", "function": { + "name": "triangle_properties_get", "strict": true, + "parameters": {"type": "object", "properties": { + "side1": {"type": "integer"}, "get_area": {"type": "boolean"}}, + "required": ["side1"], "additionalProperties": false}}}); + let tight = json!({"type": "function", "function": { + "name": "ping", "strict": true, + "parameters": {"type": "object", "properties": {"x": {"type": "integer"}}, + "required": ["x"], "additionalProperties": false}}}); + + // Upstream's contract: the strict declaration makes the loose schema a 400. + let mut beam = request(json!({"tools": [loose.clone(), tight.clone()]})); + assert!(ValidateRequest::validate(&beam).is_err()); + + // Beam route: the loose tool is served as non-strict, the compliant one keeps strict. + assert_eq!( + relax_noncompliant_strict_tools_for_beam(&mut beam, Some("beam"), Some("beam")), + 1 + ); + let tools = beam.inner.tools.as_ref().unwrap(); + assert_eq!(tools[0].function.strict, None); + assert_eq!(tools[1].function.strict, Some(true)); + assert!(ValidateRequest::validate(&beam).is_ok()); + assert_eq!( + relax_noncompliant_strict_tools_for_beam(&mut beam, Some("beam"), None), + 0 + ); + + // Every other route keeps the rejection. + let mut hermes = request(json!({"tools": [loose]})); + assert_eq!( + relax_noncompliant_strict_tools_for_beam(&mut hermes, Some("hermes"), None), + 0 + ); + assert_eq!( + hermes.inner.tools.as_ref().unwrap()[0].function.strict, + Some(true) + ); + assert!(ValidateRequest::validate(&hermes).is_err()); + } + + #[test] + fn beam_open_final_channel_holds_only_the_answer_and_forces_no_call() { + let preprocessor = preprocessor("beam", StructuralTagMode::Off); + let open_answer = PromptReasoningPrefill { + legacy: false, + unified: false, + prefix: None, + response: true, + beam_channel: Some(dynamo_parsers_v2::UnifiedParserStartingState::Response), + }; + + let mut common_request = preprocessed_request(); + preprocessor + .apply_tool_choice_guided_decoding( + &request(json!({"tools": tools()})), + &mut common_request, + open_answer, + ) + .unwrap(); + let tag = installed_structural_tag(&common_request).expect("answer rest"); + assert_eq!(tag["format"]["elements"][0]["type"], "any_text"); + assert!(!tag.to_string().contains("to=get_weather")); + + let mut common_request = preprocessed_request(); + let applied = preprocessor + .apply_tool_choice_guided_decoding( + &request(json!({"tools": tools(), "tool_choice": "required"})), + &mut common_request, + open_answer, + ) + .unwrap(); + assert_eq!(applied, GuidedToolConstraint::None); + assert!(installed_structural_tag(&common_request).is_none()); + } + #[test] fn structural_tag_off_preserves_kimi_k2_required_native_tag() { let request = request(json!({ diff --git a/lib/llm/src/protocols/anthropic/types.rs b/lib/llm/src/protocols/anthropic/types.rs index 8154824c6f12..ae015ccf9271 100644 --- a/lib/llm/src/protocols/anthropic/types.rs +++ b/lib/llm/src/protocols/anthropic/types.rs @@ -167,6 +167,7 @@ impl TryFrom for NvCreateChatCompletionRequest { .transpose()?; Ok(NvCreateChatCompletionRequest { + preserve_thinking: None, inner: dynamo_protocols::types::CreateChatCompletionRequest { messages, model: req.model, diff --git a/lib/llm/src/protocols/common.rs b/lib/llm/src/protocols/common.rs index caca44f842ca..a963b4677e27 100644 --- a/lib/llm/src/protocols/common.rs +++ b/lib/llm/src/protocols/common.rs @@ -665,6 +665,13 @@ pub struct OutputOptions { #[serde(default, skip_serializing_if = "Option::is_none")] pub no_stop_trim: Option, + /// Keep the text of the model's own end-of-turn stop tokens (the hidden system/EOS + /// ids) in the decoded output while user stop strings and user stop token ids are + /// still trimmed. The frontend sets it for grammars whose parser must see the turn + /// terminator (Beam's `<|call|>` / `<|return|>`); it is not a request field. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub retain_model_stop_token_text: Option, + /// Number of log probabilities to return per output token. /// Note that the implementation follows the OpenAI API: The return /// result includes the log probabilities on the `logprobs` most likely diff --git a/lib/llm/src/protocols/common/input_trigger.rs b/lib/llm/src/protocols/common/input_trigger.rs index 51781d5a513c..8c096fa9d39c 100644 --- a/lib/llm/src/protocols/common/input_trigger.rs +++ b/lib/llm/src/protocols/common/input_trigger.rs @@ -150,6 +150,7 @@ mod tests { messages: Vec, ) -> NvCreateChatCompletionRequest { NvCreateChatCompletionRequest { + preserve_thinking: None, inner: CreateChatCompletionRequest { messages, model: "test".to_string(), diff --git a/lib/llm/src/protocols/openai.rs b/lib/llm/src/protocols/openai.rs index 60ec42056d93..e5cc597f79e8 100644 --- a/lib/llm/src/protocols/openai.rs +++ b/lib/llm/src/protocols/openai.rs @@ -29,7 +29,7 @@ pub mod pooling; pub mod rerank; pub mod responses; pub mod stream_aggregator; -mod strict_schema; +pub(crate) mod strict_schema; pub mod tools; pub mod validate; pub mod videos; @@ -275,6 +275,7 @@ impl OutputOptionsProvider for T { let return_tokens_as_token_ids = self.get_return_tokens_as_token_ids(); Ok(common::OutputOptions { + retain_model_stop_token_text: None, logprobs, prompt_logprobs, skip_special_tokens, @@ -389,6 +390,13 @@ pub enum GuidedToolConstraint { GuidedJsonNamed { tool_name: String }, /// The model emits one call object or an array of `{name, parameters}` objects. GuidedJsonRequired, + /// The parsers crate's Beam turn grammar holds the whole turn on the model's own + /// wire: a non-empty analysis, then native calls whose arguments keep their schema + /// types, or the answer. Parsing stays native, framing every held tool's typed + /// arguments the way the grammar wrote them. `frames_answer`: the request's + /// `response_format` document is held inside the final channel behind the analysis + /// instead of being installed as a bare JSON constraint from the first token. + BeamTurnGrammar { frames_answer: bool }, } impl GuidedToolConstraint { @@ -402,6 +410,17 @@ impl GuidedToolConstraint { pub(crate) fn uses_structural_tag(&self) -> bool { matches!(self, Self::StructuralTag) } + + /// The request's `response_format` is produced inside the model's answer channel + /// behind its reasoning, so the output is a native turn and not a bare document. + pub(crate) fn frames_structured_answer(&self) -> bool { + matches!( + self, + Self::BeamTurnGrammar { + frames_answer: true + } + ) + } } #[derive(Clone, Debug, Serialize, Deserialize)] @@ -538,7 +557,14 @@ impl ParsingOptions { if !whole_response_decoder { self.tool_call_parser = None; } - self.guided_tool_constraint = GuidedToolConstraint::None; + // The Beam turn grammar shaped the whole turn (answer included) and installs + // no tool JSON, so it still describes the output when calls are not allowed. + if !matches!( + self.guided_tool_constraint, + GuidedToolConstraint::BeamTurnGrammar { .. } + ) { + self.guided_tool_constraint = GuidedToolConstraint::None; + } } self } diff --git a/lib/llm/src/protocols/openai/chat_completions.rs b/lib/llm/src/protocols/openai/chat_completions.rs index 85b2cfe144c8..c23e05955326 100644 --- a/lib/llm/src/protocols/openai/chat_completions.rs +++ b/lib/llm/src/protocols/openai/chat_completions.rs @@ -118,6 +118,13 @@ pub struct NvCreateChatCompletionRequest { #[serde(default, skip_serializing_if = "Option::is_none")] pub thinking_token_budget: Option, + /// Keep prior assistant reasoning in the rendered history (Beam serving path). + /// Beam-route request extension: supplied prior-turn reasoning is always kept in the + /// rendered prompt (`true` and omitted are the same); `false` is accepted for + /// compatibility and has no effect — omit `reasoning_content` to drop reasoning. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub preserve_thinking: Option, + /// Runtime media decoding parameters, forwarded verbatim to the worker when the /// worker owns decoding. When the frontend decodes, these override the MDC defaults. /// Kept opaque so options the frontend does not own pass through untouched. @@ -2195,3 +2202,71 @@ mod tests { assert_eq!(stop_conditions.max_thinking_tokens, None); } } + +#[cfg(test)] +mod preserve_thinking_tests { + use super::*; + use serde_json::json; + + fn chat_request_with(extra: &serde_json::Value) -> NvCreateChatCompletionRequest { + let mut body = json!({ + "model": "test-model", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 20 + }); + for (key, value) in extra.as_object().expect("fixture is an object") { + body[key] = value.clone(); + } + serde_json::from_value(body).expect("Failed to deserialize request") + } + + #[test] + fn preserve_thinking_omitted_is_none_and_not_serialized() { + let request = chat_request_with(&json!({})); + assert_eq!(request.preserve_thinking, None); + assert!(request.unsupported_fields.is_empty()); + assert!( + serde_json::to_value(&request) + .unwrap() + .get("preserve_thinking") + .is_none() + ); + ValidateRequest::validate(&request).unwrap(); + } + + #[test] + fn preserve_thinking_true_roundtrips_and_validates() { + let request = chat_request_with(&json!({"preserve_thinking": true})); + assert_eq!(request.preserve_thinking, Some(true)); + assert!(request.unsupported_fields.is_empty()); + ValidateRequest::validate(&request).unwrap(); + let value = serde_json::to_value(&request).unwrap(); + assert_eq!(value["preserve_thinking"], json!(true)); + let back: NvCreateChatCompletionRequest = serde_json::from_value(value).unwrap(); + assert_eq!(back.preserve_thinking, Some(true)); + } + + #[test] + fn preserve_thinking_false_roundtrips_and_validates() { + let request = chat_request_with(&json!({"preserve_thinking": false})); + assert_eq!(request.preserve_thinking, Some(false)); + ValidateRequest::validate(&request).unwrap(); + let value = serde_json::to_value(&request).unwrap(); + assert_eq!(value["preserve_thinking"], json!(false)); + } + + #[test] + fn preserve_thinking_rejects_non_bool_values() { + for bad in [json!("true"), json!(1), json!({})] { + let body = json!({ + "model": "test-model", + "messages": [{"role": "user", "content": "hi"}], + "preserve_thinking": bad, + }); + assert!( + serde_json::from_value::(body).is_err(), + "{bad} must not deserialize" + ); + } + } +} diff --git a/lib/llm/src/protocols/openai/chat_completions/tool_parser_v2.rs b/lib/llm/src/protocols/openai/chat_completions/tool_parser_v2.rs index c2aba138fa42..995091f9918a 100644 --- a/lib/llm/src/protocols/openai/chat_completions/tool_parser_v2.rs +++ b/lib/llm/src/protocols/openai/chat_completions/tool_parser_v2.rs @@ -294,7 +294,9 @@ pub(crate) fn parse_complete_guided_json( }) .collect() } - GuidedToolConstraint::None | GuidedToolConstraint::StructuralTag => { + GuidedToolConstraint::None + | GuidedToolConstraint::StructuralTag + | GuidedToolConstraint::BeamTurnGrammar { .. } => { anyhow::bail!("request did not install guided JSON") } } diff --git a/lib/llm/src/protocols/openai/chat_completions/unified_parser.rs b/lib/llm/src/protocols/openai/chat_completions/unified_parser.rs index 30aa97c53db1..058a962d8a6e 100644 --- a/lib/llm/src/protocols/openai/chat_completions/unified_parser.rs +++ b/lib/llm/src/protocols/openai/chat_completions/unified_parser.rs @@ -40,9 +40,9 @@ use std::sync::LazyLock; use async_stream::stream; use dynamo_parsers::tool_calling::ToolDefinition; use dynamo_parsers_v2::{ - InvalidGuidedPayloadPolicy, Tool, UnifiedEvent, UnifiedParser, UnifiedParserEvent, - UnifiedParserInit, UnifiedParserOutput, UnifiedParserStartingState, UnifiedToolOutputMode, - create_unified_parser_for_family, + ArgumentFraming, InvalidGuidedPayloadPolicy, Tool, ToolCallPolicy, UnifiedEvent, UnifiedParser, + UnifiedParserEvent, UnifiedParserInit, UnifiedParserOutput, UnifiedParserStartingState, + UnifiedToolOutputMode, create_unified_parser_for_family, }; use dynamo_protocols::types::{ ChatChoiceStream, ChatCompletionMessageContent, ChatCompletionMessageToolCall, @@ -73,6 +73,10 @@ pub(crate) const DEEPSEEK_V4_UNIFIED_FAMILY: &str = "deepseek_v4"; pub(crate) const GEMMA4_UNIFIED_FAMILY: &str = "gemma4"; pub(crate) const GLM47_UNIFIED_FAMILY: &str = "glm47"; pub(crate) const MUSE_UNIFIED_FAMILY: &str = "muse_glimmer"; +/// The `dynamo-parsers-v2` unified family that serves Beam. Dynamo has no v1 Beam +/// parser, so the `beam` + `beam` pair routes to the unified parser without +/// `DYN_PARSER_VERSION=2`, like DeepSeek V4.1. +pub(crate) const BEAM_UNIFIED_FAMILY: &str = "beam"; /// Dynamo's `--dyn-tool-call-parser` name that pairs into [`QWEN3_UNIFIED_FAMILY`]. const QWEN3_TOOL_CALL_PARSER: &str = "qwen3_coder"; @@ -155,9 +159,9 @@ fn selected_family_for_selection( ) -> Option<&'static str> { match selection { dynamo_runtime::config::ParserVersion::Auto => configured.filter(|family| { - *family == DEEPSEEK_V41_UNIFIED_FAMILY - && tool_call_parser == Some(DEEPSEEK_V41_UNIFIED_FAMILY) - && reasoning_parser == Some(DEEPSEEK_V41_UNIFIED_FAMILY) + [DEEPSEEK_V41_UNIFIED_FAMILY, BEAM_UNIFIED_FAMILY].contains(family) + && tool_call_parser == Some(*family) + && reasoning_parser == Some(*family) }), dynamo_runtime::config::ParserVersion::V1 => None, dynamo_runtime::config::ParserVersion::V2 => configured, @@ -267,6 +271,16 @@ pub(crate) struct UnifiedRequestPolicy { } impl UnifiedRequestPolicy { + /// The policy as the installed grammar leaves it: a `response_format` the grammar + /// holds inside the model's answer channel is native channel output, not the bare + /// JSON document the structured-response probes expect. + fn under(self, constraint: &GuidedToolConstraint) -> Self { + Self { + structured_response: self.structured_response && !constraint.frames_structured_answer(), + ..self + } + } + fn starting_state( self, family: &str, @@ -642,6 +656,10 @@ fn reasoning_marker_prefill( contains_unquoted_marker(content, "<|start|>") || contains_unquoted_marker(content, "<|eom|>") } + BEAM_UNIFIED_FAMILY => BEAM_OPENERS + .iter() + .chain(BEAM_CLOSERS.iter()) + .any(|marker| content.contains(marker)), _ => false, } }; @@ -687,6 +705,9 @@ pub(crate) fn stream_prefill( // so the model emits `` itself when it thinks. QWEN3_UNIFIED_FAMILY => UnifiedParserStartingState::None, DEEPSEEK_V41_UNIFIED_FAMILY => UnifiedParserStartingState::None, + // Beam's prompt ends at the `<|start|>assistant` header slot; the model opens + // its own channel and the grammar owns that open slot. + BEAM_UNIFIED_FAMILY => UnifiedParserStartingState::None, _ => UnifiedParserStartingState::None, } } @@ -778,10 +799,42 @@ fn detect_prefill(family: &str, content: &str) -> anyhow::Result Ok(infer_marker_prefill(content, "<|start|>", "<|eom|>")), + BEAM_UNIFIED_FAMILY => Ok(beam_prefill(content)), other => anyhow::bail!("no prefill detector for unified parser family '{other}'"), } } +/// Beam's channel markers are special tokens the engine emits as such, so a literal +/// never hides inside quotes: plain substring positions are the right reading. +const BEAM_OPENERS: [&str; 2] = ["<|channel|>", "<|start|>"]; +const BEAM_CLOSERS: [&str; 3] = ["<|end|>", "<|return|>", "<|call|>"]; + +/// Beam's channel state read back off complete output. +/// +/// The prompt ends at the `<|start|>assistant` header slot, so output normally opens +/// its own channel (`<|channel|>…`, or a framed `<|start|>…` message): `None`. A closer +/// before any opener means the prompt had opened a channel: `<|end|>` closes the +/// analysis channel it opened (`Reasoning`), `<|return|>` the final one (`Response`). +/// Output with no control marker at all is a final answer the prompt opened +/// (`Response`), which is also what already-decoded text looks like when the +/// aggregator re-reads it. +fn beam_prefill(content: &str) -> UnifiedParserStartingState { + let opener = BEAM_OPENERS + .iter() + .filter_map(|marker| content.find(marker)) + .min(); + let closer = BEAM_CLOSERS + .iter() + .filter_map(|marker| content.find(marker).map(|at| (at, *marker))) + .min_by_key(|(at, _)| *at); + match (opener, closer) { + (Some(open), Some((close, _))) if open < close => UnifiedParserStartingState::None, + (Some(_), None) => UnifiedParserStartingState::None, + (_, Some((_, "<|end|>"))) => UnifiedParserStartingState::Reasoning, + (_, Some(_)) | (None, None) => UnifiedParserStartingState::Response, + } +} + fn reasoning_marker_pair(family: &str) -> Option<(&'static str, &'static str)> { match family { "qwen3" | "deepseek_v4" | "deepseek_v41" | "glm47" | "kimi_k2" => { @@ -1054,10 +1107,38 @@ fn tool_output_mode(constraint: &GuidedToolConstraint) -> UnifiedToolOutputMode GuidedToolConstraint::GuidedJsonRequired => { UnifiedToolOutputMode::GuidedJson { named_tool: None } } - GuidedToolConstraint::None | GuidedToolConstraint::StructuralTag => { - UnifiedToolOutputMode::Native - } + GuidedToolConstraint::None + | GuidedToolConstraint::StructuralTag + | GuidedToolConstraint::BeamTurnGrammar { .. } => UnifiedToolOutputMode::Native, + } +} + +/// How the native parser frames argument values for this request: the Beam turn +/// grammar writes every held tool's non-string parameters as JSON, so those slots own +/// their strings and a closing tag inside one is data (see +/// `dynamo_parsers_v2::ArgumentFraming::SchemaConstrained`); without a grammar a value +/// runs to the first closing tag. +fn argument_framing(constraint: &GuidedToolConstraint) -> ArgumentFraming { + match constraint { + GuidedToolConstraint::BeamTurnGrammar { .. } => ArgumentFraming::SchemaConstrained, + GuidedToolConstraint::None + | GuidedToolConstraint::StructuralTag + | GuidedToolConstraint::GuidedJsonNamed { .. } + | GuidedToolConstraint::GuidedJsonRequired => ArgumentFraming::Lenient, + } +} + +fn apply_argument_framing( + parser: &mut dyn UnifiedParser, + argument_framing: ArgumentFraming, +) -> anyhow::Result<()> { + if argument_framing != ArgumentFraming::default() { + parser.set_tool_call_policy(ToolCallPolicy { + malformed: Default::default(), + argument_framing, + })?; } + Ok(()) } /// The malformed-guided-payload contract for one request, selected by the @@ -1175,6 +1256,21 @@ fn unified_parser_control_markers(family: &str) -> Vec { .map(str::to_string), ); } + if family == BEAM_UNIFIED_FAMILY { + markers.extend( + [ + "<|start|>", + "<|channel|>", + "<|message|>", + "<|constrain|>", + "<|end|>", + "<|return|>", + "<|call|>", + ] + .into_iter() + .map(str::to_string), + ); + } if let Some(marker) = native_family_marker(family) { markers.push(marker.to_string()); } @@ -1195,6 +1291,7 @@ fn native_family_marker(family: &str) -> Option<&'static str> { "kimi_k2" => Some("<|tool_call_begin|>"), "kimi_k3" => Some("<|open|>call"), "muse_glimmer" => Some(" Some("<|start|>assistant to="), _ => None, } } @@ -1642,6 +1739,7 @@ impl ChoiceState { prefill: impl Into, tool_output_mode: UnifiedToolOutputMode, guided_streaming: bool, + argument_framing: ArgumentFraming, ) -> anyhow::Result { let prefill = prefill.into(); let mut parser = create_unified_parser_for_family(family, tools)?; @@ -1667,6 +1765,7 @@ impl ChoiceState { tool_output_mode: tool_output_mode.clone(), invalid_guided_payload: invalid_guided_payload_policy(guided_streaming), })?; + apply_argument_framing(parser.as_mut(), argument_framing)?; if let Some(prefix) = prefill.prefix { // Replay only consumed structural bytes; the grammar retains ownership // of an optional channel label that generation has not completed yet. @@ -2280,6 +2379,7 @@ fn text_delta(text: String) -> Vec { } /// The aggregated result of parsing one complete (non-streaming) output. +#[derive(Debug)] pub(crate) struct CompleteOutput { pub incomplete_native_call: bool, pub text: String, @@ -2319,6 +2419,7 @@ pub(crate) fn parse_complete_with_policy( tool_definitions: &[ToolDefinition], policy: UnifiedRequestPolicy, ) -> anyhow::Result { + let policy = policy.under(guided_tool_constraint); // A remote frontend may reconstruct guided JSON even when the backend returned // native reasoning/content without a tool call. Keep that recovery under the // same grammar owner as ordinary native batch output. @@ -2395,6 +2496,7 @@ fn parse_complete_in_mode( tool_output_mode: effective_mode, ..UnifiedParserInit::default() })?; + apply_argument_framing(parser.as_mut(), argument_framing(guided_tool_constraint))?; let mut text = String::new(); let mut reasoning = String::new(); @@ -2965,7 +3067,9 @@ where starting_state: prefill, prefix: prefill_prefix, } = prefill.into(); + let policy = policy.under(&guided_tool_constraint); let prefill = policy.starting_state(family, prefill); + let argument_framing = argument_framing(&guided_tool_constraint); let tools = to_v2_tools(tool_definitions.as_deref()); stream! { let mut undecided_guided_prefix: HashMap = HashMap::new(); @@ -3460,7 +3564,7 @@ where starting_state: choice_prefill, prefix: (choice_prefill == UnifiedParserStartingState::Reasoning).then_some(prefill_prefix).flatten(), }; - match ChoiceState::new(family, &tools, choice_prefill, mode, guided_streaming) { + match ChoiceState::new(family, &tools, choice_prefill, mode, guided_streaming, argument_framing) { Ok(mut state) => { // A raw run resuming after an already-parsed detour // gets a brand-new parser instance; seed it with any @@ -4085,6 +4189,12 @@ pub(crate) mod tests { for native in [ GuidedToolConstraint::None, GuidedToolConstraint::StructuralTag, + GuidedToolConstraint::BeamTurnGrammar { + frames_answer: false, + }, + GuidedToolConstraint::BeamTurnGrammar { + frames_answer: true, + }, ] { assert_eq!( tool_output_mode(&native), @@ -4094,6 +4204,54 @@ pub(crate) mod tests { } } + #[test] + fn the_beam_turn_grammar_frames_arguments_and_owns_a_structured_answer() { + // Only the grammar that typed the arguments frames them. + for (constraint, framing) in [ + (GuidedToolConstraint::None, ArgumentFraming::Lenient), + ( + GuidedToolConstraint::StructuralTag, + ArgumentFraming::Lenient, + ), + ( + GuidedToolConstraint::GuidedJsonRequired, + ArgumentFraming::Lenient, + ), + ( + GuidedToolConstraint::BeamTurnGrammar { + frames_answer: false, + }, + ArgumentFraming::SchemaConstrained, + ), + ] { + assert_eq!(argument_framing(&constraint), framing, "{constraint:?}"); + } + // A `response_format` the grammar holds in the final channel is native output. + let structured = UnifiedRequestPolicy { + reasoning_disabled: false, + structured_response: true, + }; + assert!( + structured + .under(&GuidedToolConstraint::BeamTurnGrammar { + frames_answer: false, + }) + .structured_response + ); + assert!( + !structured + .under(&GuidedToolConstraint::BeamTurnGrammar { + frames_answer: true, + }) + .structured_response + ); + assert!( + structured + .under(&GuidedToolConstraint::None) + .structured_response + ); + } + #[test] fn batch_native_marker_overrides_a_reconstructed_guided_constraint() { for constraint in [ @@ -4299,6 +4457,7 @@ pub(crate) mod tests { UnifiedParserStartingState::None, UnifiedToolOutputMode::Native, true, + ArgumentFraming::default(), ) .expect("qwen3 unified parser") } @@ -8665,6 +8824,7 @@ pub(crate) mod tests { UnifiedParserStartingState::Response, mode.clone(), true, + ArgumentFraming::default(), ) .unwrap(); let mut events = state.push(prefix); @@ -8725,6 +8885,7 @@ pub(crate) mod tests { UnifiedParserStartingState::Response, mode.clone(), true, + ArgumentFraming::default(), ) .unwrap(); let mut split_events = split_state.push(&short_raw[..split]); @@ -10012,6 +10173,7 @@ pub(crate) mod tests { UnifiedParserStartingState::None, mode, true, + ArgumentFraming::default(), ) .unwrap(); let mut library = Vec::new(); @@ -10201,6 +10363,7 @@ pub(crate) mod tests { UnifiedParserStartingState::None, UnifiedToolOutputMode::Native, true, + ArgumentFraming::default(), ) .unwrap(); let events = state.push(&input[..stop]); @@ -12173,6 +12336,7 @@ pub(crate) mod tests { UnifiedParserStartingState::None, UnifiedToolOutputMode::Native, true, + ArgumentFraming::default(), ) .unwrap(); let started = std::time::Instant::now(); @@ -12218,6 +12382,7 @@ pub(crate) mod tests { UnifiedParserStartingState::None, UnifiedToolOutputMode::Native, true, + ArgumentFraming::default(), ) .unwrap(); state.push(&input[..split]); @@ -12234,3 +12399,227 @@ pub(crate) mod tests { } } } + +#[cfg(test)] +mod beam_routing_tests { + use super::*; + use dynamo_runtime::config::ParserVersion; + + /// The turn grammar writes a held tool's non-string parameters as JSON, so a closing + /// tag inside one of their strings is data; and it holds a `response_format` answer + /// in the final channel behind the analysis, so that output is a native turn. + #[test] + fn beam_turn_grammar_frames_held_json_arguments_and_the_structured_answer() { + let tools = vec![ToolDefinition { + name: "render".into(), + parameters: Some(serde_json::json!({ + "type": "object", + "properties": { + "doc": {"type": "object", "properties": {"html": {"type": "string"}}}, + "title": {"type": "string"} + }, + "required": ["doc", "title"] + })), + strict: None, + }]; + let content = concat!( + "<|channel|>analysis<|message|>Rendering.<|end|>", + "<|start|>assistant to=render<|channel|>analysis <|constrain|>xml<|message|>", + "{\"html\":\"x inside\"}", + "t<|call|>" + ); + let expected = serde_json::json!({ + "doc": {"html": "x inside"}, + "title": "t" + }); + let held = GuidedToolConstraint::BeamTurnGrammar { + frames_answer: false, + }; + let batch = parse_complete_with_policy( + BEAM_UNIFIED_FAMILY, + content, + &held, + &tools, + UnifiedRequestPolicy::default(), + ) + .unwrap(); + assert_eq!(batch.reasoning, "Rendering."); + assert_eq!(batch.tool_calls.len(), 1, "{batch:?}"); + assert_eq!(batch.tool_calls[0].function.name, "render"); + assert_eq!( + serde_json::from_str::(&batch.tool_calls[0].function.arguments) + .unwrap(), + expected + ); + // Without the grammar a value runs to the first closing tag. + let lenient = parse_complete_with_policy( + BEAM_UNIFIED_FAMILY, + content, + &GuidedToolConstraint::None, + &tools, + UnifiedRequestPolicy::default(), + ) + .unwrap(); + assert!( + lenient.tool_calls.first().is_none_or(|call| { + serde_json::from_str::(&call.function.arguments).ok() + != Some(expected.clone()) + }), + "{lenient:?}" + ); + + // Streaming frames the same way at every split. + let v2_tools = to_v2_tools(Some(&tools)); + for split in content + .char_indices() + .map(|(at, _)| at) + .chain(std::iter::once(content.len())) + { + let mut state = ChoiceState::new( + BEAM_UNIFIED_FAMILY, + &v2_tools, + UnifiedParserStartingState::None, + UnifiedToolOutputMode::Native, + true, + argument_framing(&held), + ) + .unwrap(); + let mut wire = Vec::new(); + for events in [ + state.push(&content[..split]), + state.push(&content[split..]), + state.finish(), + ] { + wire.extend(state.choices_for(&empty_choice(0), events, true, None)); + } + let calls: Vec<_> = wire + .iter() + .flat_map(|choice| choice.delta.tool_calls.iter().flatten()) + .collect(); + let arguments: String = calls + .iter() + .filter_map(|call| call.function.as_ref()?.arguments.as_deref()) + .collect(); + assert_eq!( + serde_json::from_str::(&arguments).unwrap(), + expected, + "split {split}" + ); + let reasoning: String = wire + .iter() + .filter_map(|choice| choice.delta.reasoning_content.as_deref()) + .collect(); + assert_eq!(reasoning, "Rendering.", "split {split}"); + } + + // A `response_format` answer the grammar framed keeps the reasoning and the + // channel structure, and is not probed as a bare JSON document. + let answer = concat!( + "<|channel|>analysis<|message|>Thinking.<|end|>", + "<|start|>assistant<|channel|>final<|message|>{\"ok\":true}<|return|>" + ); + let structured = UnifiedRequestPolicy { + reasoning_disabled: false, + structured_response: true, + }; + let framed = parse_complete_with_policy( + BEAM_UNIFIED_FAMILY, + answer, + &GuidedToolConstraint::BeamTurnGrammar { + frames_answer: true, + }, + &[], + structured, + ) + .unwrap(); + assert_eq!(framed.reasoning, "Thinking."); + assert_eq!(framed.text, "{\"ok\":true}"); + assert!(framed.tool_calls.is_empty()); + } + + #[test] + fn beam_pair_routes_to_the_unified_parser_without_an_explicit_generation() { + assert_eq!( + configured_family(Some("beam"), Some("beam")), + Some(BEAM_UNIFIED_FAMILY) + ); + assert_eq!( + selected_family_for_selection( + Some(BEAM_UNIFIED_FAMILY), + ParserVersion::Auto, + Some("beam"), + Some("beam"), + ), + Some(BEAM_UNIFIED_FAMILY) + ); + // Half a pair keeps the split path: there is no v1 Beam parser to fall back to, + // so the operator names both halves. + assert_eq!( + selected_family_for_selection( + Some(BEAM_UNIFIED_FAMILY), + ParserVersion::Auto, + Some("beam"), + None, + ), + None + ); + assert_eq!( + selected_family_for_selection( + Some(BEAM_UNIFIED_FAMILY), + ParserVersion::V1, + Some("beam"), + Some("beam"), + ), + None + ); + assert!(FAMILY_NAMES.contains(&BEAM_UNIFIED_FAMILY)); + } + + #[test] + fn beam_prefill_reads_the_channel_the_prompt_opened() { + assert_eq!( + beam_prefill( + "<|channel|>analysis<|message|>t<|end|><|start|>assistant<|channel|>final<|message|>a<|return|>" + ), + UnifiedParserStartingState::None + ); + assert_eq!( + beam_prefill( + "<|start|>assistant to=get_weather<|channel|>analysis <|constrain|>xml<|message|><|call|>" + ), + UnifiedParserStartingState::None + ); + assert_eq!( + beam_prefill("thought<|end|><|start|>assistant<|channel|>final<|message|>a<|return|>"), + UnifiedParserStartingState::Reasoning + ); + assert_eq!( + beam_prefill("answer<|return|>"), + UnifiedParserStartingState::Response + ); + assert_eq!( + beam_prefill("plain decoded answer"), + UnifiedParserStartingState::Response + ); + assert_eq!( + detect_prefill("beam", "plain").unwrap(), + UnifiedParserStartingState::Response + ); + } + + #[test] + fn beam_control_markers_cover_the_whole_grammar() { + let markers = unified_parser_control_markers("beam"); + for marker in [ + "<|start|>", + "<|channel|>", + "<|message|>", + "<|constrain|>", + "<|end|>", + "<|return|>", + "<|call|>", + ] { + assert!(markers.iter().any(|m| m == marker), "{marker}"); + } + } +} diff --git a/lib/llm/src/protocols/openai/completions.rs b/lib/llm/src/protocols/openai/completions.rs index f53c425c0727..19471027eda1 100644 --- a/lib/llm/src/protocols/openai/completions.rs +++ b/lib/llm/src/protocols/openai/completions.rs @@ -71,6 +71,83 @@ pub struct NvCreateCompletionResponse { pub inner: dynamo_protocols::types::CreateCompletionResponse, #[serde(skip_serializing_if = "Option::is_none")] pub nvext: Option, + /// The prompt a choice echoes back (`echo: true`), built once the engine reports + /// its prompt logprobs. In-process only, never serialized on its own: the unary + /// aggregator prepends it to the choice and the SSE path folds it into the chunk + /// that carries it (see [`NvCreateCompletionResponse::fold_echo`]). + #[serde(skip)] + pub echo: Option, +} + +/// The prompt side of a completions choice for `echo: true`: the prompt text and, when +/// the request asked for `logprobs`, one entry per prompt token in the same layout as +/// the completion's `logprobs` (position 0 has no logprob, as in the OpenAI API). +/// +/// The two parts can travel on different chunks: the text goes out on the choice's +/// first chunk, the prompt logprobs on the chunk that carries the engine's payload +/// (the final one on vLLM and SGLang). A part that already went out is empty here. +#[derive(Debug, Clone, PartialEq)] +pub struct CompletionEcho { + pub choice_index: u32, + pub text: String, + pub logprobs: Option, +} + +impl CompletionEcho { + /// Combine the parts of one choice's echo that arrived on different chunks. + pub fn merge(&mut self, other: CompletionEcho) { + if self.text.is_empty() { + self.text = other.text; + } + if self.logprobs.is_none() { + self.logprobs = other.logprobs; + } + } + + /// Put the echoed prompt ahead of the completion: text first, then the prompt + /// positions before the generated ones. Generated `text_offset`s move past the + /// echoed text. + pub fn prepend_to(self, choice: &mut dynamo_protocols::types::Choice) { + let echoed_len = self.text.len() as u32; + choice.text.insert_str(0, &self.text); + let Some(mut logprobs) = self.logprobs else { + return; + }; + if let Some(generated) = choice.logprobs.take() { + logprobs.tokens.extend(generated.tokens); + logprobs.token_logprobs.extend(generated.token_logprobs); + logprobs.top_logprobs.extend(generated.top_logprobs); + logprobs.text_offset.extend( + generated + .text_offset + .into_iter() + .map(|offset| offset + echoed_len), + ); + } + choice.logprobs = Some(logprobs); + } +} + +impl NvCreateCompletionResponse { + /// Fold a pending echo into its choice. For consumers that serialize this chunk as + /// is (SSE); the unary aggregator prepends the echo once the stream has ended. + pub fn fold_echo(&mut self) { + let Some(echo) = self.echo.take() else { + return; + }; + match self + .inner + .choices + .iter_mut() + .find(|choice| choice.index == echo.choice_index) + { + Some(choice) => echo.prepend_to(choice), + None => tracing::warn!( + choice_index = echo.choice_index, + "completions echo dropped: the chunk carrying it has no matching choice" + ), + } + } } /// Schema-only import for the wholly re-exported async-openai response type. @@ -362,7 +439,11 @@ impl ResponseFactory { system_fingerprint: self.system_fingerprint.clone(), usage, }; - NvCreateCompletionResponse { inner, nvext: None } + NvCreateCompletionResponse { + inner, + nvext: None, + echo: None, + } } } @@ -466,9 +547,12 @@ impl OpenAIOutputOptionsProvider for NvCreateCompletionRequest { fn get_prompt_logprobs(&self) -> Option { self.common.prompt_logprobs.or_else(|| { + // `echo` rebuilds the prompt positions from the engine's prompt logprobs, + // with as many alternatives per position as `logprobs` asks for. self.inner .echo - .and_then(|echo| if echo { Some(1) } else { None }) + .filter(|echo| *echo) + .map(|_| self.inner.logprobs.map_or(1, u32::from)) }) } diff --git a/lib/llm/src/protocols/openai/completions/aggregator.rs b/lib/llm/src/protocols/openai/completions/aggregator.rs index 4928131d44b7..73412c7824ea 100644 --- a/lib/llm/src/protocols/openai/completions/aggregator.rs +++ b/lib/llm/src/protocols/openai/completions/aggregator.rs @@ -6,7 +6,7 @@ use std::collections::HashMap; use dynamo_runtime::error::DynamoError; use futures::{Stream, StreamExt, TryStreamExt}; -use super::NvCreateCompletionResponse; +use super::{CompletionEcho, NvCreateCompletionResponse}; use crate::protocols::{ Annotated, DataStream, codec::{Message, SseCodecError}, @@ -31,6 +31,20 @@ struct DeltaChoice { text: String, finish_reason: Option, logprobs: Option, + /// The echoed prompt (`echo: true`), prepended once the stream has ended. + echo: Option, +} + +impl DeltaChoice { + fn empty(index: u32) -> Self { + Self { + index, + text: String::new(), + finish_reason: None, + logprobs: None, + echo: None, + } + } } impl Default for DeltaAggregator { @@ -75,19 +89,24 @@ impl DeltaAggregator { aggregator.system_fingerprint = Some(system_fingerprint); } merge_response_nvext(&mut aggregator.nvext, delta.nvext); + if let Some(echo) = delta.echo { + let index = echo.choice_index; + let state_choice = aggregator + .choices + .entry(index) + .or_insert_with(|| DeltaChoice::empty(index)); + match state_choice.echo.as_mut() { + Some(existing) => existing.merge(echo), + None => state_choice.echo = Some(echo), + } + } // handle the choices for choice in delta.inner.choices { - let state_choice = - aggregator - .choices - .entry(choice.index) - .or_insert(DeltaChoice { - index: choice.index, - text: "".to_string(), - finish_reason: None, - logprobs: None, - }); + let state_choice = aggregator + .choices + .entry(choice.index) + .or_insert_with(|| DeltaChoice::empty(choice.index)); state_choice.text.push_str(&choice.text); @@ -152,6 +171,7 @@ impl DeltaAggregator { let response = NvCreateCompletionResponse { inner, nvext: aggregator.nvext, + echo: None, }; Ok(response) @@ -162,12 +182,16 @@ impl From for dynamo_protocols::types::Choice { fn from(delta: DeltaChoice) -> Self { let finish_reason = delta.finish_reason.map(Into::into); - dynamo_protocols::types::Choice { + let mut choice = dynamo_protocols::types::Choice { index: delta.index, text: delta.text, finish_reason, logprobs: delta.logprobs, + }; + if let Some(echo) = delta.echo { + echo.prepend_to(&mut choice); } + choice } } @@ -239,7 +263,11 @@ mod tests { object: "text_completion".to_string(), }; - let response = NvCreateCompletionResponse { inner, nvext: None }; + let response = NvCreateCompletionResponse { + inner, + nvext: None, + echo: None, + }; Annotated { data: Some(response), @@ -398,7 +426,11 @@ mod tests { object: "text_completion".to_string(), }; - let response = NvCreateCompletionResponse { inner, nvext: None }; + let response = NvCreateCompletionResponse { + inner, + nvext: None, + echo: None, + }; let annotated_delta = Annotated { data: Some(response), @@ -474,4 +506,43 @@ mod tests { ); assert_eq!(error.message(), "unsupported logprobs"); } + + #[tokio::test] + async fn echo_is_prepended_to_the_aggregated_choice() { + // The text rides the first chunk, the prompt logprobs the final one. + let mut first = create_test_delta(0, "Hi", None, Some(-0.3)); + first.data.as_mut().expect("delta data").echo = Some(CompletionEcho { + choice_index: 0, + text: "Prompt".to_string(), + logprobs: None, + }); + let mut last = create_test_delta(0, "!", Some("stop".to_string()), Some(-0.7)); + last.data.as_mut().expect("delta data").echo = Some(CompletionEcho { + choice_index: 0, + text: String::new(), + logprobs: Some(dynamo_protocols::types::Logprobs { + tokens: vec!["Prompt".to_string()], + token_logprobs: vec![None], + top_logprobs: vec![serde_json::Value::Null], + text_offset: vec![], + }), + }); + + let response = DeltaAggregator::apply( + futures::stream::iter(vec![first, last]), + ParsingOptions::default(), + ) + .await + .expect("aggregation"); + + assert!(response.echo.is_none()); + let choice = &response.inner.choices[0]; + assert_eq!(choice.text, "PromptHi!"); + let logprobs = choice.logprobs.as_ref().expect("logprobs"); + assert_eq!(logprobs.tokens, vec!["Prompt", "Hi", "!"]); + assert_eq!(logprobs.token_logprobs, vec![None, Some(-0.3), Some(-0.7)]); + assert_eq!(logprobs.top_logprobs[0], serde_json::Value::Null); + // Generated offsets (0 in `create_test_delta`) move past the echoed text. + assert_eq!(logprobs.text_offset, vec![6, 6]); + } } diff --git a/lib/llm/src/protocols/openai/completions/delta.rs b/lib/llm/src/protocols/openai/completions/delta.rs index d417fbeb074a..39fc8f3c42a3 100644 --- a/lib/llm/src/protocols/openai/completions/delta.rs +++ b/lib/llm/src/protocols/openai/completions/delta.rs @@ -1,14 +1,16 @@ // SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +use std::collections::{HashMap, HashSet}; use std::sync::Arc; -use super::{NvCreateCompletionRequest, NvCreateCompletionResponse}; +use super::{CompletionEcho, NvCreateCompletionRequest, NvCreateCompletionResponse}; use crate::{ protocols::{ common::{ self, extensions::{NvExtProvider, NvExtResponseInput}, + llm_backend::PromptLogprobs, timing::RequestTracker, }, openai::{ @@ -16,6 +18,7 @@ use crate::{ delta_common::{self, DeltaGeneratorOptions, DeltaGeneratorState}, }, }, + tokenizers::traits::Tokenizer, types::TokenIdType, }; @@ -42,6 +45,122 @@ impl NvCreateCompletionRequest { pub struct DeltaGenerator { state: DeltaGeneratorState, + echo: Option, +} + +/// What `echo: true` needs to rebuild the prompt side of a choice: the prompt the +/// preprocessor sent to the engine and a tokenizer for its per-token text. +struct EchoState { + prompt_token_ids: Arc>, + /// The client's prompt verbatim when it was text; token prompts are decoded. + prompt_text: Option, + tokenizer: Arc, + skip_special_tokens: bool, + /// Choices whose prompt text already went out (once per choice). + text_emitted: HashSet, + /// Choices whose prompt logprobs already went out, or were given up on. + logprobs_emitted: HashSet, + /// Generated text and logprobs held back per choice while its prompt logprobs are + /// still to come, so the prompt positions precede the generated ones on the wire. + held: HashMap)>, +} + +fn append_logprobs( + into: &mut Option, + more: dynamo_protocols::types::Logprobs, +) { + match into { + None => *into = Some(more), + Some(existing) => { + existing.tokens.extend(more.tokens); + existing.token_logprobs.extend(more.token_logprobs); + existing.top_logprobs.extend(more.top_logprobs); + existing.text_offset.extend(more.text_offset); + } + } +} + +impl EchoState { + fn token_text(&self, token_id: TokenIdType, return_as_ids: bool) -> String { + if return_as_ids { + return format!("token_id:{token_id}"); + } + // Single-token decode on purpose: BPE merge and leading-space rules differ + // between single-token and sequence decode (see `Backend`'s top-logprob fill). + self.tokenizer + .decode(&[token_id], self.skip_special_tokens) + .map(String::from) + .unwrap_or_default() + } + + fn prompt_text(&self) -> String { + if let Some(text) = &self.prompt_text { + return text.clone(); + } + self.tokenizer + .decode(&self.prompt_token_ids, self.skip_special_tokens) + .map(String::from) + .unwrap_or_default() + } + + /// One `logprobs` entry per prompt token from the engine's per-position maps, + /// in the completions layout: the token's own logprob (none at position 0) and a + /// token-to-logprob object of the alternatives the engine reported. + fn logprobs( + &self, + prompt_logprobs: &PromptLogprobs, + return_as_ids: bool, + ) -> dynamo_protocols::types::Logprobs { + if prompt_logprobs.len() != self.prompt_token_ids.len() { + tracing::warn!( + positions = prompt_logprobs.len(), + prompt_tokens = self.prompt_token_ids.len(), + "completions echo: the engine's prompt logprobs do not cover the prompt one-to-one" + ); + } + let n = self.prompt_token_ids.len(); + let mut tokens = Vec::with_capacity(n); + let mut token_logprobs = Vec::with_capacity(n); + let mut top_logprobs = Vec::with_capacity(n); + for (position, &token_id) in self.prompt_token_ids.iter().enumerate() { + let text = self.token_text(token_id, return_as_ids); + let entries = prompt_logprobs.get(position).and_then(Option::as_ref); + let selected = entries + .and_then(|entries| entries.get(&token_id)) + .map(|entry| entry.logprob); + top_logprobs.push(match entries { + Some(entries) => { + let mut values = serde_json::Map::new(); + for (&alternative, entry) in entries { + let key = if return_as_ids { + format!("token_id:{alternative}") + } else { + entry + .decoded_token + .clone() + .unwrap_or_else(|| self.token_text(alternative, false)) + }; + values.insert(key, serde_json::json!(entry.logprob)); + } + // Distinct ids can decode to the same string; the prompt token's + // own probability wins, as on the completion side. + if let Some(logprob) = selected { + values.insert(text.clone(), serde_json::json!(logprob)); + } + serde_json::Value::Object(values) + } + None => serde_json::Value::Null, + }); + token_logprobs.push(selected); + tokens.push(text); + } + dynamo_protocols::types::Logprobs { + tokens, + token_logprobs, + top_logprobs, + text_offset: vec![], + } + } } impl DeltaGenerator { @@ -53,9 +172,35 @@ impl DeltaGenerator { model, options, ), + echo: None, } } + /// Echo the prompt back on every choice (`echo: true`). `prompt_text` is the + /// client's prompt when it was text; token prompts are decoded with `tokenizer`. + /// The prompt text rides the first chunk of each choice, so a streaming client + /// sees it before any generated text. When `logprobs` was requested, the prompt + /// logprobs ride the chunk that carries the engine's payload (vLLM and SGLang + /// attach it to the final chunk), and the choice's generated text and logprobs + /// are held back until then, so the wire order stays prompt, then completion. + pub fn enable_echo( + &mut self, + prompt_token_ids: Arc>, + prompt_text: Option, + tokenizer: Arc, + skip_special_tokens: bool, + ) { + self.echo = Some(EchoState { + prompt_token_ids, + prompt_text, + tokenizer, + skip_special_tokens, + text_emitted: HashSet::new(), + logprobs_emitted: HashSet::new(), + held: HashMap::new(), + }); + } + /// Returns the request tracker. Tracking is always enabled. For sharing with PreprocessedRequest. pub fn tracker(&self) -> Arc { self.state.tracker() @@ -165,7 +310,11 @@ impl DeltaGenerator { }, }; - NvCreateCompletionResponse { inner, nvext: None } + NvCreateCompletionResponse { + inner, + nvext: None, + echo: None, + } } /// Creates a final usage-only chunk for OpenAI compliance. @@ -186,7 +335,11 @@ impl DeltaGenerator { usage: Some(usage), }; - NvCreateCompletionResponse { inner, nvext: None } + NvCreateCompletionResponse { + inner, + nvext: None, + echo: None, + } } /// Check if usage tracking is enabled @@ -259,12 +412,86 @@ impl crate::protocols::openai::DeltaGeneratorExt for self.state.tracker_ref().record_finish(); } + // The engine's prompt logprobs have two readers: the nvext field when the + // client selected it, and `echo`, which rebuilds the prompt side of the choice + // from them. + let enable_logprobs = self.state.options().enable_logprobs; + let (wants_echo_text, wants_echo_logprobs) = match &self.echo { + Some(echo) => ( + !echo.text_emitted.contains(&index), + enable_logprobs && !echo.logprobs_emitted.contains(&index), + ), + None => (false, false), + }; + let prompt_logprobs = + if self.state.options().response_fields.prompt_logprobs || wants_echo_logprobs { + common::llm_backend::prompt_logprobs_from_engine_data(delta.engine_data.as_ref())? + } else { + None + }; + if wants_echo_logprobs + && let Some(echo) = self.echo.as_mut() + && let Some(choice) = response + .inner + .choices + .iter_mut() + .find(|choice| choice.index == index) + { + if prompt_logprobs.is_none() && finish_reason.is_none() { + // Hold this chunk's generated side until the prompt logprobs arrive. + let held = echo.held.entry(index).or_default(); + held.0.push_str(&std::mem::take(&mut choice.text)); + if let Some(logprobs) = choice.logprobs.take() { + append_logprobs(&mut held.1, logprobs); + } + } else if let Some((text, logprobs)) = echo.held.remove(&index) { + // Release what was held ahead of this chunk's own generated side. + choice.text.insert_str(0, &text); + let mut released = logprobs; + if let Some(current) = choice.logprobs.take() { + append_logprobs(&mut released, current); + } + choice.logprobs = released; + } + } + if wants_echo_text || wants_echo_logprobs { + let return_as_ids = self.state.options().return_tokens_as_token_ids; + let echo = self.echo.as_mut().expect("checked above"); + let text = if wants_echo_text { + echo.text_emitted.insert(index); + echo.prompt_text() + } else { + String::new() + }; + let logprobs = match &prompt_logprobs { + Some(payload) if wants_echo_logprobs => { + echo.logprobs_emitted.insert(index); + Some(echo.logprobs(payload, return_as_ids)) + } + None if wants_echo_logprobs && finish_reason.is_some() => { + tracing::warn!( + "completions echo: the engine reported no prompt logprobs; the prompt was echoed as text only" + ); + echo.logprobs_emitted.insert(index); + None + } + _ => None, + }; + if !text.is_empty() || logprobs.is_some() { + response.echo = Some(CompletionEcho { + choice_index: index, + text, + logprobs, + }); + } + } + // Build the nvext response payload via the shared gating helper on // `NvExtResponseFieldSelection` (see `nvext.rs`). Both chat and // completions delta generators go through the same helper so the gating // rules stay in one place. let prompt_logprobs_payload = if self.state.options().response_fields.prompt_logprobs { - common::llm_backend::prompt_logprobs_from_engine_data(delta.engine_data.as_ref())? + prompt_logprobs } else { None }; @@ -911,4 +1138,212 @@ mod tests { ); } } + + struct MapTokenizer; + + impl crate::tokenizers::traits::Encoder for MapTokenizer { + fn encode(&self, _input: &str) -> anyhow::Result { + anyhow::bail!("echo tests never encode") + } + + fn encode_batch( + &self, + _inputs: &[&str], + ) -> anyhow::Result> { + anyhow::bail!("echo tests never encode") + } + } + + impl crate::tokenizers::traits::Decoder for MapTokenizer { + fn decode( + &self, + token_ids: &[TokenIdType], + skip_special_tokens: bool, + ) -> anyhow::Result { + let text: String = token_ids + .iter() + .map(|id| match id { + 1 => "", + 5 => "Hello", + 7 => " world", + 9 => " there", + _ => "?", + }) + .filter(|piece| !(skip_special_tokens && *piece == "")) + .collect(); + Ok(crate::tokenizers::traits::DecodeResult::Complete(text)) + } + } + + impl crate::tokenizers::traits::Tokenizer for MapTokenizer {} + + fn echo_request(logprobs: Option) -> NvCreateCompletionRequest { + let mut request = create_test_request(); + request.inner.echo = Some(true); + request.inner.logprobs = logprobs; + request + } + + /// A chunk with one generated token and the engine's prompt logprobs for the + /// prompt `[1, 5, 7]`: nothing at position 0, two candidates at position 1. + fn prompt_logprobs_output(finished: bool) -> BackendOutput { + let mut output = final_backend_output(); + if !finished { + output.finish_reason = None; + } + output.log_probs = Some(vec![-0.1]); + output.engine_data = Some(serde_json::json!({ + "prompt_logprobs": [ + null, + {"5": {"logprob": -0.25}, "9": {"logprob": -1.5, "decoded_token": " there"}}, + {"7": {"logprob": -0.5}} + ] + })); + output + } + + #[test] + fn echo_rebuilds_the_prompt_positions_from_the_engine_prompt_logprobs() { + let request = echo_request(Some(2)); + let mut generator = request.response_generator("req-echo".to_string()); + generator.enable_echo(Arc::new(vec![1, 5, 7]), None, Arc::new(MapTokenizer), false); + + let mut response = generator + .choice_from_postprocessor(prompt_logprobs_output(false)) + .expect("choice generation"); + let echo = response + .echo + .clone() + .expect("the chunk carrying the prompt logprobs echoes the prompt"); + assert_eq!(echo.choice_index, 0); + assert_eq!(echo.text, "Hello world"); + let logprobs = echo.logprobs.expect("logprobs were requested"); + assert_eq!(logprobs.tokens, vec!["", "Hello", " world"]); + assert_eq!(logprobs.token_logprobs, vec![None, Some(-0.25), Some(-0.5)]); + assert_eq!(logprobs.top_logprobs[0], serde_json::Value::Null); + assert_eq!( + logprobs.top_logprobs[1], + serde_json::json!({"Hello": -0.25, " there": -1.5}) + ); + assert_eq!( + logprobs.top_logprobs[2], + serde_json::json!({" world": -0.5}) + ); + assert_eq!( + serde_json::to_value(&response) + .expect("serialize chunk") + .get("echo"), + None, + "the echo must not leak as a field of its own" + ); + + response.fold_echo(); + assert!(response.echo.is_none()); + let choice = &response.inner.choices[0]; + assert_eq!(choice.text, "Hello worldhello"); + let folded = choice + .logprobs + .as_ref() + .expect("prompt positions ahead of the completion's"); + assert_eq!(folded.tokens, vec!["", "Hello", " world", "hello"]); + assert_eq!( + folded.token_logprobs, + vec![None, Some(-0.25), Some(-0.5), Some(-0.1)] + ); + + let next = generator + .choice_from_postprocessor(prompt_logprobs_output(true)) + .expect("choice generation"); + assert!(next.echo.is_none(), "a choice echoes its prompt once"); + } + + #[test] + fn echo_falls_back_to_the_prompt_text_when_the_engine_reports_no_prompt_logprobs() { + let request = echo_request(Some(1)); + let mut generator = request.response_generator("req-echo-text".to_string()); + generator.enable_echo( + Arc::new(vec![5, 7]), + Some("Hello world".to_string()), + Arc::new(MapTokenizer), + true, + ); + + let mut partial = final_backend_output(); + partial.finish_reason = None; + partial.engine_data = None; + let response = generator + .choice_from_postprocessor(partial) + .expect("choice generation"); + let echo = response + .echo + .expect("the choice's first chunk echoes the prompt, ahead of any generated text"); + assert_eq!(echo.text, "Hello world", "a text prompt is echoed verbatim"); + assert!(echo.logprobs.is_none(), "no prompt logprobs yet"); + + let mut last = final_backend_output(); + last.engine_data = None; + let response = generator + .choice_from_postprocessor(last) + .expect("choice generation"); + assert!( + response.echo.is_none(), + "the engine never reported prompt logprobs: nothing more to echo" + ); + } + + #[test] + fn echo_logprobs_ride_the_final_chunk_after_the_text() { + let request = echo_request(Some(2)); + let mut generator = request.response_generator("req-echo-split".to_string()); + generator.enable_echo(Arc::new(vec![1, 5, 7]), None, Arc::new(MapTokenizer), false); + + // First chunk: no prompt logprobs yet (vLLM and SGLang attach them to the + // final chunk), so only the text goes out. + let mut partial = final_backend_output(); + partial.finish_reason = None; + partial.engine_data = None; + let response = generator + .choice_from_postprocessor(partial) + .expect("choice generation"); + let echo = response + .echo + .expect("the first chunk echoes the prompt text"); + assert_eq!(echo.text, "Hello world"); + assert!(echo.logprobs.is_none()); + assert_eq!( + response.inner.choices[0].text, "", + "the generated side is held until the prompt logprobs arrive" + ); + + // Final chunk: the prompt logprobs, the held text and this chunk's own, in order. + let mut response = generator + .choice_from_postprocessor(prompt_logprobs_output(true)) + .expect("choice generation"); + let echo = response + .echo + .clone() + .expect("the chunk carrying the prompt logprobs echoes them"); + assert_eq!(echo.text, "", "the prompt text already went out"); + let logprobs = echo.logprobs.expect("logprobs were requested"); + assert_eq!(logprobs.tokens, vec!["", "Hello", " world"]); + response.fold_echo(); + let choice = &response.inner.choices[0]; + assert_eq!( + choice.text, "hellohello", + "held text first, then this chunk's" + ); + let folded = choice.logprobs.as_ref().expect("prompt positions first"); + assert_eq!(folded.tokens, vec!["", "Hello", " world", "hello"]); + assert_eq!(folded.token_logprobs.last(), Some(&Some(-0.1))); + } + + #[test] + fn echo_is_not_built_unless_requested() { + let request = create_test_request(); + let mut generator = request.response_generator("req-no-echo".to_string()); + let response = generator + .choice_from_postprocessor(prompt_logprobs_output(true)) + .expect("choice generation"); + assert!(response.echo.is_none()); + } } diff --git a/lib/llm/src/protocols/openai/responses/mod.rs b/lib/llm/src/protocols/openai/responses/mod.rs index 0df64658135c..927b4701142c 100644 --- a/lib/llm/src/protocols/openai/responses/mod.rs +++ b/lib/llm/src/protocols/openai/responses/mod.rs @@ -1026,6 +1026,7 @@ impl TryFrom for NvCreateChatCompletionRequest { .transpose()?; Ok(NvCreateChatCompletionRequest { + preserve_thinking: None, inner: CreateChatCompletionRequest { messages, model: resp.inner.model.unwrap_or_default(), diff --git a/lib/llm/src/protocols/openai/strict_schema.rs b/lib/llm/src/protocols/openai/strict_schema.rs index 6dfa4268e74f..9caa2ec1d3b2 100644 --- a/lib/llm/src/protocols/openai/strict_schema.rs +++ b/lib/llm/src/protocols/openai/strict_schema.rs @@ -31,7 +31,7 @@ const TYPES: &[&str] = &[ "null", "boolean", "object", "array", "number", "integer", "string", ]; -pub(super) fn validate_strict_function(function: &FunctionObject) -> Result<()> { +pub(crate) fn validate_strict_function(function: &FunctionObject) -> Result<()> { if function.strict != Some(true) { return Ok(()); } diff --git a/lib/llm/src/protocols/openai/validate.rs b/lib/llm/src/protocols/openai/validate.rs index 60cbf3595683..67b533df58da 100644 --- a/lib/llm/src/protocols/openai/validate.rs +++ b/lib/llm/src/protocols/openai/validate.rs @@ -67,8 +67,10 @@ pub const MAX_TOP_LOGPROBS: u8 = 20; /// Minimum allowed value for `logprobs` in completion requests pub const MIN_LOGPROBS: u8 = 0; -/// Maximum allowed value for `logprobs` in completion requests -pub const MAX_LOGPROBS: u8 = 5; +/// Maximum allowed value for `logprobs` in completion requests. The same bound as +/// chat's `top_logprobs`: engines score up to 20 alternatives per position, and +/// `echo` + `logprobs` teacher-forcing clients ask for that many on the prompt. +pub const MAX_LOGPROBS: u8 = 20; /// Minimum allowed value for `n` (number of choices) pub const MIN_N: u8 = 1; diff --git a/lib/llm/src/request_trace/agent_context.rs b/lib/llm/src/request_trace/agent_context.rs index e4e9508051e9..54c9de2d6e80 100644 --- a/lib/llm/src/request_trace/agent_context.rs +++ b/lib/llm/src/request_trace/agent_context.rs @@ -782,6 +782,7 @@ mod tests { usage: None, }, nvext: None, + echo: None, })]; for response in &responses { record_completion_finish_reason_metadata(&trace_state.finish_reason_metadata, response); diff --git a/lib/llm/tests/parallel_tool_call_integration.rs b/lib/llm/tests/parallel_tool_call_integration.rs index dfd0310492a8..76367b23237a 100644 --- a/lib/llm/tests/parallel_tool_call_integration.rs +++ b/lib/llm/tests/parallel_tool_call_integration.rs @@ -96,6 +96,7 @@ fn create_mock_chat_completion_request() -> NvCreateChatCompletionRequest { media_io_kwargs: None, return_tokens_as_token_ids: None, thinking_token_budget: None, + preserve_thinking: None, unsupported_fields: Default::default(), } } diff --git a/lib/llm/tests/preprocessor.rs b/lib/llm/tests/preprocessor.rs index 29faf27d913d..198eca316c79 100644 --- a/lib/llm/tests/preprocessor.rs +++ b/lib/llm/tests/preprocessor.rs @@ -265,6 +265,7 @@ impl Request { media_io_kwargs: None, return_tokens_as_token_ids: None, thinking_token_budget: None, + preserve_thinking: None, unsupported_fields: Default::default(), } } @@ -954,6 +955,7 @@ mod context_length_validation { media_io_kwargs: None, return_tokens_as_token_ids: None, thinking_token_budget: None, + preserve_thinking: None, unsupported_fields: Default::default(), } } @@ -1335,6 +1337,7 @@ mod embedding_without_chat_template { return_tokens_as_token_ids: None, unsupported_fields: Default::default(), thinking_token_budget: None, + preserve_thinking: None, } } diff --git a/lib/llm/tests/test_common_ext.rs b/lib/llm/tests/test_common_ext.rs index 8b0c5ddc6968..68a37b8eb4f5 100644 --- a/lib/llm/tests/test_common_ext.rs +++ b/lib/llm/tests/test_common_ext.rs @@ -72,6 +72,7 @@ fn test_sampling_parameters_include_stop_str_in_output_extraction() { media_io_kwargs: None, return_tokens_as_token_ids: None, thinking_token_budget: None, + preserve_thinking: None, unsupported_fields: Default::default(), }; @@ -317,6 +318,7 @@ fn test_completions_common_values() { fn test_serialization_preserves_structure() { // Test that serialization preserves the flattened structure let request = NvCreateChatCompletionRequest { + preserve_thinking: None, inner: dynamo_protocols::types::CreateChatCompletionRequest { model: "test-model".to_string(), messages: vec![dynamo_protocols::types::ChatCompletionRequestMessage::User( @@ -398,6 +400,7 @@ fn test_sampling_parameters_extraction() { media_io_kwargs: None, return_tokens_as_token_ids: None, thinking_token_budget: None, + preserve_thinking: None, unsupported_fields: Default::default(), }; diff --git a/lib/llm/tests/test_stop_behavior.rs b/lib/llm/tests/test_stop_behavior.rs index 606076f630e7..7615d3ec40a8 100644 --- a/lib/llm/tests/test_stop_behavior.rs +++ b/lib/llm/tests/test_stop_behavior.rs @@ -88,6 +88,7 @@ fn make_decoder( stop_conditions, include_stop_str, false, + false, None, None, ) @@ -183,7 +184,15 @@ fn user_stop_token_reports_distinct_trigger() { stop_token_ids_hidden: Some(vec![EOS]), ..Default::default() }; - let mut decoder = Decoder::new(decode_stream, stop_conditions, false, false, None, None); + let mut decoder = Decoder::new( + decode_stream, + stop_conditions, + false, + false, + false, + None, + None, + ); let result = decoder.process_token_ids(&[HI, STOP]).unwrap(); assert_eq!(result.text.as_deref(), Some("hi")); @@ -283,7 +292,15 @@ fn visible_stop_token_flushes_and_orders_prior_jailed_prefix() { stop: Some(vec!["hiya".to_string()]), ..Default::default() }; - let mut decoder = Decoder::new(decode_stream, stop_conditions, false, false, None, None); + let mut decoder = Decoder::new( + decode_stream, + stop_conditions, + false, + false, + false, + None, + None, + ); let result = decoder.process_token_ids(&[HI, STOP]).unwrap(); assert_eq!(result.text.as_deref(), Some("hiSTOP")); @@ -371,6 +388,7 @@ fn no_stop_trim_preserves_seed_order_and_special_token_policy() { }, false, no_stop_trim, + false, None, Some("o".into()), ); diff --git a/lib/llm/tests/test_streaming_usage.rs b/lib/llm/tests/test_streaming_usage.rs index 5d75857e7438..7e086cbe0fe3 100644 --- a/lib/llm/tests/test_streaming_usage.rs +++ b/lib/llm/tests/test_streaming_usage.rs @@ -204,6 +204,7 @@ fn create_chat_request( }; NvCreateChatCompletionRequest { + preserve_thinking: None, inner, common: Default::default(), nvext: None, @@ -798,6 +799,7 @@ fn create_nonstreaming_chat_request() -> NvCreateChatCompletionRequest { }; NvCreateChatCompletionRequest { + preserve_thinking: None, inner, common: Default::default(), nvext: None, diff --git a/lib/llm/tests/tool_choice.rs b/lib/llm/tests/tool_choice.rs index b4225fe26706..dbbdb3a1bc9e 100644 --- a/lib/llm/tests/tool_choice.rs +++ b/lib/llm/tests/tool_choice.rs @@ -59,6 +59,7 @@ fn create_test_request() -> NvCreateChatCompletionRequest { )]; NvCreateChatCompletionRequest { + preserve_thinking: None, inner: CreateChatCompletionRequest { model: "test-model".to_string(), messages, diff --git a/lib/llm/tests/tool_choice_finish_reasons.rs b/lib/llm/tests/tool_choice_finish_reasons.rs index f6dc184dda66..27c38bf1e851 100644 --- a/lib/llm/tests/tool_choice_finish_reasons.rs +++ b/lib/llm/tests/tool_choice_finish_reasons.rs @@ -22,6 +22,7 @@ fn create_test_request() -> NvCreateChatCompletionRequest { )]; NvCreateChatCompletionRequest { + preserve_thinking: None, inner: CreateChatCompletionRequest { model: "test-model".to_string(), messages, diff --git a/tests/report_pytest_markers.py b/tests/report_pytest_markers.py index 677d052c5983..695922141a0d 100755 --- a/tests/report_pytest_markers.py +++ b/tests/report_pytest_markers.py @@ -253,6 +253,9 @@ "vllm.reasoning.mistral_reasoning_parser", "vllm.reasoning.qwen3_engine_reasoning_parser", "vllm.renderers", + "vllm.renderers.chat_utils", + "vllm.renderers.chat_utils.chat_template", + "vllm.renderers.chat_utils.tool_call", "vllm.renderers.embed_utils", "vllm.sampling_params", "vllm.tokenizers",