diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index 46cc34cea210..b65882eded07 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -1103,6 +1103,10 @@ class ChatCompletionRequest(OpenAIBaseModel): # emits per turn, so parallel emission remains model behavior either way. parallel_tool_calls: Optional[bool] = None user: Optional[str] = None + # Standard OpenAI field, accepted for compatibility and otherwise ignored: + # KV cache reuse needs no client hint, and cache_salt is the field that + # isolates cached prefixes between users. + prompt_cache_key: Optional[str] = None reasoning_effort: Optional[ReasoningEffort | Literal[ "low", "medium", "high", "max", "none"]] = Field( default=ReasoningEffort.LOW, diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index b7bb5eeb30ad..85eee2eed5b5 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -208,6 +208,7 @@ l0_cpu: - unittest/llmapi/apps/test_openai_chat_disagg_multimodal.py - unittest/llmapi/apps/test_openai_chat_multimodal_prompt.py - unittest/llmapi/apps/test_openai_protocol_chat_store.py + - unittest/llmapi/apps/test_openai_protocol_prompt_cache_key.py - unittest/llmapi/apps/test_openai_server_iteration_stats.py - unittest/llmapi/apps/test_reasoning_prompt_resolution.py - unittest/llmapi/apps/test_render_callers.py diff --git a/tests/unittest/api_stability/references/trtllm_serve_api.yaml b/tests/unittest/api_stability/references/trtllm_serve_api.yaml index 3f2568cf1394..4d8b4b458312 100644 --- a/tests/unittest/api_stability/references/trtllm_serve_api.yaml +++ b/tests/unittest/api_stability/references/trtllm_serve_api.yaml @@ -646,6 +646,12 @@ models: default: null status: stable required: false + prompt_cache_key: + kind: openai + type: Optional[str] + default: null + status: stable + required: false reasoning_effort: kind: openai type: Optional[ReasoningEffort | Literal['low', 'medium', 'high', 'max', 'none']] diff --git a/tests/unittest/llmapi/apps/test_openai_protocol_prompt_cache_key.py b/tests/unittest/llmapi/apps/test_openai_protocol_prompt_cache_key.py new file mode 100644 index 000000000000..22a4c43ba618 --- /dev/null +++ b/tests/unittest/llmapi/apps/test_openai_protocol_prompt_cache_key.py @@ -0,0 +1,73 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""ChatCompletionRequest must accept the standard `prompt_cache_key` field. + +The request schema is extra="forbid", so the field must exist for requests +from OpenAI clients that set it (for example agent frameworks that pass a +per-session key) to validate. The value is not used. +""" + +import pytest +from pydantic import ValidationError + +from tensorrt_llm.serve.openai_protocol import ChatCompletionRequest + +pytestmark = pytest.mark.cpu_only + + +def _request(**extra) -> ChatCompletionRequest: + return ChatCompletionRequest( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + **extra, + ) + + +@pytest.mark.parametrize("value", ["session-1234", ""]) +def test_prompt_cache_key_accepted(value): + req = _request(prompt_cache_key=value) + assert req.prompt_cache_key == value + + +def test_prompt_cache_key_defaults_to_none(): + assert _request().prompt_cache_key is None + + +def test_prompt_cache_key_accepted_from_json(): + req = ChatCompletionRequest.model_validate_json( + '{"model": "test-model", "messages": [{"role": "user", "content": "hi"}],' + ' "prompt_cache_key": "session-1234"}' + ) + assert req.prompt_cache_key == "session-1234" + + +@pytest.mark.parametrize("value", [1, ["session"], {"key": "session"}]) +def test_prompt_cache_key_must_be_a_string(value): + with pytest.raises(ValidationError): + _request(prompt_cache_key=value) + + +def test_prompt_cache_key_leaves_other_fields_unchanged(): + plain = _request(cache_salt="tenant") + keyed = _request(cache_salt="tenant", prompt_cache_key="session-1234") + assert keyed.cache_salt == "tenant" + assert keyed.model_dump(exclude={"prompt_cache_key"}) == plain.model_dump( + exclude={"prompt_cache_key"} + ) + + +def test_unknown_fields_remain_rejected(): + with pytest.raises(ValidationError): + _request(prompt_cache_keys="typo")