diff --git a/integrations/nvidia/src/haystack_integrations/components/rankers/nvidia/ranker.py b/integrations/nvidia/src/haystack_integrations/components/rankers/nvidia/ranker.py index 0105697fc5..7022e05dbd 100644 --- a/integrations/nvidia/src/haystack_integrations/components/rankers/nvidia/ranker.py +++ b/integrations/nvidia/src/haystack_integrations/components/rankers/nvidia/ranker.py @@ -3,18 +3,15 @@ # SPDX-License-Identifier: Apache-2.0 import os -import warnings from dataclasses import replace from typing import Any -from haystack import Document, component, default_from_dict, default_to_dict, logging +from haystack import Document, component, default_from_dict, default_to_dict from haystack.utils import Secret, deserialize_secrets_inplace from haystack_integrations.components.rankers.nvidia.truncate import RankerTruncateMode from haystack_integrations.utils.nvidia import DEFAULT_API_URL, Client, NimBackend, is_hosted, url_validation -logger = logging.getLogger(__name__) - @component class NvidiaRanker: @@ -84,6 +81,8 @@ def __init__( :param timeout: Timeout for request calls, if not set it is inferred from the `NVIDIA_TIMEOUT` environment variable or set to 60 by default. + + :raises ValueError: If `top_k` is not > 0. """ if model is not None and not isinstance(model, str): msg = "Ranker expects the `model` parameter to be a string." @@ -96,6 +95,9 @@ def __init__( if not isinstance(top_k, int): msg = "Ranker expects the `top_k` parameter to be an integer." raise TypeError(msg) + if top_k <= 0: + msg = f"top_k must be > 0, but got {top_k}" + raise ValueError(msg) # todo: detect default in non-hosted case (when api_url is provided) self.truncate = truncate @@ -202,6 +204,7 @@ def run(self, query: str, documents: list[Document], top_k: int | None = None) - :param top_k: The number of documents to return. :raises TypeError: If the arguments are of the wrong type. + :raises ValueError: If `top_k` is not > 0. :returns: A dictionary containing the ranked documents. """ @@ -224,11 +227,10 @@ def run(self, query: str, documents: list[Document], top_k: int | None = None) - if len(documents) == 0: return {"documents": []} + if top_k is not None and top_k <= 0: + msg = f"top_k must be > 0, but got {top_k}" + raise ValueError(msg) top_k = top_k if top_k is not None else self.top_k - if top_k < 1: - logger.warning("top_k should be at least 1, returning nothing") - warnings.warn("top_k should be at least 1, returning nothing", stacklevel=2) - return {"documents": []} assert self.backend is not None diff --git a/integrations/nvidia/tests/test_ranker.py b/integrations/nvidia/tests/test_ranker.py index 2988e5b63d..e03b384dab 100644 --- a/integrations/nvidia/tests/test_ranker.py +++ b/integrations/nvidia/tests/test_ranker.py @@ -221,20 +221,27 @@ def test_nim_integration(self): assert len(response) == 2 assert {response[0].content, response[1].content} == {documents[0].content, documents[1].content} - def test_top_k_warn(self, monkeypatch) -> None: + @pytest.mark.parametrize("top_k", [0, -1]) + def test_top_k_init_invalid(self, monkeypatch, top_k: int) -> None: monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key") + with pytest.raises(ValueError, match=rf"top_k must be > 0, but got {top_k}"): + NvidiaRanker(top_k=top_k) - client = NvidiaRanker(top_k=0) + @pytest.mark.parametrize("top_k", [0, -1]) + def test_top_k_run_invalid(self, monkeypatch, top_k: int) -> None: + monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key") + client = NvidiaRanker() client.warm_up() - with pytest.warns(UserWarning) as record0: - client.run("query", [Document(content="doc")]) - assert "top_k should be at least 1" in str(record0[0].message) + with pytest.raises(ValueError, match=rf"top_k must be > 0, but got {top_k}"): + client.run("query", [Document(content="doc")], top_k=top_k) - client = NvidiaRanker(top_k=1) + def test_top_k_zero_at_run_does_not_fall_back_to_instance_top_k(self, monkeypatch) -> None: + # a runtime top_k=0 must raise, not silently fall back to the instance top_k + monkeypatch.setenv("NVIDIA_API_KEY", "fake-api-key") + client = NvidiaRanker(top_k=5) client.warm_up() - with pytest.warns(UserWarning) as record1: + with pytest.raises(ValueError, match=r"top_k must be > 0, but got 0"): client.run("query", [Document(content="doc")], top_k=0) - assert "top_k should be at least 1" in str(record1[0].message) def test_model_typeerror(self) -> None: with pytest.raises(TypeError) as e: