diff --git a/tensorrt_llm/serve/responses_utils.py b/tensorrt_llm/serve/responses_utils.py index edd36f8624f9..9c2aa65db867 100644 --- a/tensorrt_llm/serve/responses_utils.py +++ b/tensorrt_llm/serve/responses_utils.py @@ -257,9 +257,6 @@ async def store_response(self, conversation_id = self.response_to_conversation[prev_resp_id] self.conversations[conversation_id].extend(resp_msgs) - while len(self.conversations[conversation_id] - ) > self.conversation_capacity: - self._pop_conversation(resp_id) else: conversation_id = _random_uuid() self.conversations[conversation_id] = resp_msgs @@ -269,6 +266,7 @@ async def store_response(self, self.response_to_conversation[resp_id] = conversation_id self.conversation_to_response[conversation_id] = resp_id + self._trim_conversation(conversation_id) self._update_visited_conversation(conversation_id) async def pop_response(self, resp_id: Optional[str] = None) -> bool: @@ -307,12 +305,10 @@ async def store_messages(self, resp_id: str, _responses_debug_log( f" * storing at conversation: {conversation_id}") self.conversations[conversation_id] = msgs - if len(self.conversations[conversation_id] - ) > self.conversation_capacity: - self._pop_conversation(resp_id) self.response_to_conversation[resp_id] = conversation_id self.conversation_to_response[conversation_id] = resp_id + self._trim_conversation(conversation_id) self._update_visited_conversation(conversation_id) async def get_conversation_history( @@ -372,8 +368,20 @@ def _pop_conversation(self, resp_id) -> None: if conversation_id is None: return - conversation = self.conversations[conversation_id] - if len(conversation) == 0: + self._pop_conversation_by_conversation_id(conversation_id) + + def _trim_conversation(self, conversation_id: str) -> None: + conversation = self.conversations.get(conversation_id) + if conversation is None: + return + + while len(conversation) > self.conversation_capacity: + self._pop_conversation_by_conversation_id(conversation_id) + + def _pop_conversation_by_conversation_id(self, + conversation_id: str) -> None: + conversation = self.conversations.get(conversation_id) + if conversation is None or len(conversation) == 0: return is_harmony_conversation = isinstance(conversation[0], Message) diff --git a/tests/unittest/llmapi/test_responses_utils.py b/tests/unittest/llmapi/test_responses_utils.py new file mode 100644 index 000000000000..7b1a413cea82 --- /dev/null +++ b/tests/unittest/llmapi/test_responses_utils.py @@ -0,0 +1,89 @@ +# 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. + +from types import SimpleNamespace + +import pytest + +from tensorrt_llm.serve.responses_utils import ConversationHistoryStore + + +def _turns(count: int): + messages = [] + for index in range(count): + messages.extend( + [ + { + "role": "user", + "content": f"user {index}", + }, + { + "role": "assistant", + "content": f"assistant {index}", + }, + ] + ) + return messages + + +@pytest.mark.asyncio +async def test_store_response_trims_pre_stored_request_conversation(): + store = ConversationHistoryStore(resp_capacity=1) + + await store.store_messages("resp_1", _turns(3), prev_resp_id=None) + await store.store_response( + SimpleNamespace(id="resp_1"), + [ + { + "role": "assistant", + "content": "final", + } + ], + prev_resp_id=None, + ) + + conversation = await store.get_conversation_history("resp_1") + + assert len(conversation) <= store.conversation_capacity + + +@pytest.mark.asyncio +async def test_store_response_trims_previous_response_conversation(monkeypatch): + store = ConversationHistoryStore(resp_capacity=1) + + await store.store_messages("resp_prev", _turns(2), prev_resp_id=None) + + def fail_if_unmapped_response_id_is_used(_): + raise AssertionError("conversation trim used an unmapped response id") + + monkeypatch.setattr(store, "_pop_conversation", fail_if_unmapped_response_id_is_used) + + await store.store_response( + SimpleNamespace(id="resp_next"), + [ + { + "role": "assistant", + "content": "next", + } + ], + prev_resp_id="resp_prev", + ) + + conversation = await store.get_conversation_history("resp_next") + + assert len(conversation) <= store.conversation_capacity + assert ( + store.response_to_conversation["resp_next"] == store.response_to_conversation["resp_prev"] + )