From d1a5589089a6648366eb54166634497c1752bb99 Mon Sep 17 00:00:00 2001 From: Andrew White Date: Fri, 24 Jul 2026 09:28:48 -0500 Subject: [PATCH] fix: use min_output_len value in query_llm inputs query_llm filled the min_output_len input with max_output_len, so the minimum length request was ignored and the maximum length was used for both. Updates the existing test to verify min_output_len and max_output_len are passed as distinct values. Signed-off-by: Andrew White --- nemo_deploy/llm/query_llm.py | 2 +- tests/unit_tests/deploy/test_query_llm.py | 5 ++++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/nemo_deploy/llm/query_llm.py b/nemo_deploy/llm/query_llm.py index e900e4377..04e99e9a2 100755 --- a/nemo_deploy/llm/query_llm.py +++ b/nemo_deploy/llm/query_llm.py @@ -362,7 +362,7 @@ def query_llm( inputs = {"prompts": prompts} if min_output_len is not None: - inputs["min_output_len"] = np.full(prompts.shape, max_output_len, dtype=np.int_) + inputs["min_output_len"] = np.full(prompts.shape, min_output_len, dtype=np.int_) if max_output_len is not None: inputs["max_output_len"] = np.full(prompts.shape, max_output_len, dtype=np.int_) diff --git a/tests/unit_tests/deploy/test_query_llm.py b/tests/unit_tests/deploy/test_query_llm.py index e7fdfae48..9994f8442 100755 --- a/tests/unit_tests/deploy/test_query_llm.py +++ b/tests/unit_tests/deploy/test_query_llm.py @@ -501,11 +501,14 @@ def test_query_llm_with_min_output_len(self, mock_client, query): mock_instance.infer_batch.return_value = {"outputs": np.array([b"test response"])} mock_instance.model_config.outputs = [MagicMock(dtype=np.bytes_)] - # Test query with min_output_len + # Test query with min_output_len distinct from max_output_len response = query.query_llm(prompts=["test prompt"], min_output_len=10, max_output_len=100) assert isinstance(response[0], str) assert response[0] == "test response" + call_kwargs = mock_instance.infer_batch.call_args.kwargs + assert call_kwargs["min_output_len"][0] == 10 + assert call_kwargs["max_output_len"][0] == 100 @patch("nemo_deploy.llm.query_llm.ModelClient") def test_query_llm_with_logits_output(self, mock_client, query):