Add default model for NVIDIA HayStack local NIM endpoints (#915)

* initial embedder code * default model code * docs: update model docstring * tests: add userwarning * docs: literal lint fix * review changes * remove pydantic dependency * move backend, nim_backend under utils * move is_hosted to warm_up * test cases, docstring fix * error message updation Co-authored-by: Madeesh Kannan <[email protected]> * move is_hosted code to util * remove backend code * update import for is_hosted * remove util and move code to utils * fix api key issue for failing test cases * Update integrations/nvidia/tests/conftest.py --------- Co-authored-by: Madeesh Kannan <[email protected]>
deepset-ai · Oct 2, 2024 · 70ba8f6 · 70ba8f6
1 parent 6fbba9f
commit 70ba8f6
Show file tree

Hide file tree

Showing 15 changed files with 358 additions and 172 deletions.
diff --git a/integrations/nvidia/pyproject.toml b/integrations/nvidia/pyproject.toml
@@ -42,7 +42,7 @@ root = "../.."
 git_describe_command = 'git describe --tags --match="integrations/nvidia-v[0-9]*"'
 
 [tool.hatch.envs.default]
-dependencies = ["coverage[toml]>=6.5", "pytest", "pytest-rerunfailures", "haystack-pydoc-tools"]
+dependencies = ["coverage[toml]>=6.5", "pytest", "pytest-rerunfailures", "haystack-pydoc-tools", "requests_mock"]
 [tool.hatch.envs.default.scripts]
 test = "pytest --reruns 3 --reruns-delay 30 -x {args:tests}"
 test-cov = "coverage run -m pytest --reruns 3 --reruns-delay 30 -x {args:tests}"
@@ -147,6 +147,8 @@ module = [
   "haystack_integrations.*",
   "pytest.*",
   "numpy.*",
+  "requests_mock.*",
+  "pydantic.*"
 ]
 ignore_missing_imports = true
 

diff --git a/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/_nim_backend.py b/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/_nim_backend.py
diff --git a/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/backend.py b/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/backend.py
diff --git a/...rations/nvidia/src/haystack_integrations/components/embedders/nvidia/document_embedder.py b/...rations/nvidia/src/haystack_integrations/components/embedders/nvidia/document_embedder.py
@@ -1,12 +1,11 @@
+import warnings
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 from haystack import Document, component, default_from_dict, default_to_dict
 from haystack.utils import Secret, deserialize_secrets_inplace
-from haystack_integrations.utils.nvidia import url_validation
+from haystack_integrations.utils.nvidia import NimBackend, is_hosted, url_validation
 from tqdm import tqdm
 
-from ._nim_backend import NimBackend
-from .backend import EmbedderBackend
 from .truncate import EmbeddingTruncateMode
 
 _DEFAULT_API_URL = "https://ai.api.nvidia.com/v1/retrieval/nvidia"
@@ -34,7 +33,7 @@ class NvidiaDocumentEmbedder:
 
     def __init__(
         self,
-        model: str = "NV-Embed-QA",
+        model: Optional[str] = None,
         api_key: Optional[Secret] = Secret.from_env_var("NVIDIA_API_KEY"),
         api_url: str = _DEFAULT_API_URL,
         prefix: str = "",
@@ -50,6 +49,8 @@ def __init__(
 
         :param model:
             Embedding model to use.
+            If no specific model along with locally hosted API URL is provided,
+            the system defaults to the available model found using /models API.
         :param api_key:
             API key for the NVIDIA NIM.
         :param api_url:
@@ -87,9 +88,31 @@ def __init__(
             truncate = EmbeddingTruncateMode.from_str(truncate)
         self.truncate = truncate
 
-        self.backend: Optional[EmbedderBackend] = None
+        self.backend: Optional[Any] = None
         self._initialized = False
 
+        if is_hosted(api_url) and not self.model:  # manually set default model
+            self.model = "NV-Embed-QA"
+
+    def default_model(self):
+        """Set default model in local NIM mode."""
+        valid_models = [
+            model.id for model in self.backend.models() if not model.base_model or model.base_model == model.id
+        ]
+        name = next(iter(valid_models), None)
+        if name:
+            warnings.warn(
+                f"Default model is set as: {name}. \n"
+                "Set model using model parameter. \n"
+                "To get available models use available_models property.",
+                UserWarning,
+                stacklevel=2,
+            )
+            self.model = self.backend.model = name
+        else:
+            error_message = "No locally hosted model was found."
+            raise ValueError(error_message)
+
     def warm_up(self):
         """
         Initializes the component.
@@ -109,6 +132,9 @@ def warm_up(self):
 
         self._initialized = True
 
+        if not self.model:
+            self.default_model()
+
     def to_dict(self) -> Dict[str, Any]:
         """
         Serializes the component to a dictionary.

diff --git a/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/text_embedder.py b/integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/text_embedder.py
@@ -1,11 +1,10 @@
+import warnings
 from typing import Any, Dict, List, Optional, Union
 
 from haystack import component, default_from_dict, default_to_dict
 from haystack.utils import Secret, deserialize_secrets_inplace
-from haystack_integrations.utils.nvidia import url_validation
+from haystack_integrations.utils.nvidia import NimBackend, is_hosted, url_validation
 
-from ._nim_backend import NimBackend
-from .backend import EmbedderBackend
 from .truncate import EmbeddingTruncateMode
 
 _DEFAULT_API_URL = "https://ai.api.nvidia.com/v1/retrieval/nvidia"
@@ -35,7 +34,7 @@ class NvidiaTextEmbedder:
 
     def __init__(
         self,
-        model: str = "NV-Embed-QA",
+        model: Optional[str] = None,
         api_key: Optional[Secret] = Secret.from_env_var("NVIDIA_API_KEY"),
         api_url: str = _DEFAULT_API_URL,
         prefix: str = "",
@@ -47,6 +46,8 @@ def __init__(
 
         :param model:
             Embedding model to use.
+            If no specific model along with locally hosted API URL is provided,
+            the system defaults to the available model found using /models API.
         :param api_key:
             API key for the NVIDIA NIM.
         :param api_url:
@@ -71,9 +72,31 @@ def __init__(
             truncate = EmbeddingTruncateMode.from_str(truncate)
         self.truncate = truncate
 
-        self.backend: Optional[EmbedderBackend] = None
+        self.backend: Optional[Any] = None
         self._initialized = False
 
+        if is_hosted(api_url) and not self.model:  # manually set default model
+            self.model = "NV-Embed-QA"
+
+    def default_model(self):
+        """Set default model in local NIM mode."""
+        valid_models = [
+            model.id for model in self.backend.models() if not model.base_model or model.base_model == model.id
+        ]
+        name = next(iter(valid_models), None)
+        if name:
+            warnings.warn(
+                f"Default model is set as: {name}. \n"
+                "Set model using model parameter. \n"
+                "To get available models use available_models property.",
+                UserWarning,
+                stacklevel=2,
+            )
+            self.model = self.backend.model = name
+        else:
+            error_message = "No locally hosted model was found."
+            raise ValueError(error_message)
+
     def warm_up(self):
         """
         Initializes the component.
@@ -93,6 +116,9 @@ def warm_up(self):
 
         self._initialized = True
 
+        if not self.model:
+            self.default_model()
+
     def to_dict(self) -> Dict[str, Any]:
         """
         Serializes the component to a dictionary.

diff --git a/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/backend.py b/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/backend.py
diff --git a/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/generator.py b/integrations/nvidia/src/haystack_integrations/components/generators/nvidia/generator.py
@@ -1,14 +1,12 @@
 # SPDX-FileCopyrightText: 2024-present deepset GmbH <[email protected]>
 #
 # SPDX-License-Identifier: Apache-2.0
+import warnings
 from typing import Any, Dict, List, Optional
 
 from haystack import component, default_from_dict, default_to_dict
 from haystack.utils.auth import Secret, deserialize_secrets_inplace
-from haystack_integrations.utils.nvidia import url_validation
-
-from ._nim_backend import NimBackend
-from .backend import GeneratorBackend
+from haystack_integrations.utils.nvidia import NimBackend, is_hosted, url_validation
 
 _DEFAULT_API_URL = "https://integrate.api.nvidia.com/v1"
 
@@ -45,7 +43,7 @@ class NvidiaGenerator:
 
     def __init__(
         self,
-        model: str,
+        model: Optional[str] = None,
         api_url: str = _DEFAULT_API_URL,
         api_key: Optional[Secret] = Secret.from_env_var("NVIDIA_API_KEY"),
         model_arguments: Optional[Dict[str, Any]] = None,
@@ -55,6 +53,10 @@ def __init__(
 
         :param model:
             Name of the model to use for text generation.
+            See the [NVIDIA NIMs](https://ai.nvidia.com)
+            for more information on the supported models.
+            `Note`: If no specific model along with locally hosted API URL is provided,
+            the system defaults to the available model found using /models API.
             Check supported models at [NVIDIA NIM](https://ai.nvidia.com).
         :param api_key:
             API key for the NVIDIA NIM. Set it as the `NVIDIA_API_KEY` environment
@@ -72,7 +74,28 @@ def __init__(
         self._api_key = api_key
         self._model_arguments = model_arguments or {}
 
-        self._backend: Optional[GeneratorBackend] = None
+        self._backend: Optional[Any] = None
+
+        self.is_hosted = is_hosted(api_url)
+
+    def default_model(self):
+        """Set default model in local NIM mode."""
+        valid_models = [
+            model.id for model in self._backend.models() if not model.base_model or model.base_model == model.id
+        ]
+        name = next(iter(valid_models), None)
+        if name:
+            warnings.warn(
+                f"Default model is set as: {name}. \n"
+                "Set model using model parameter. \n"
+                "To get available models use available_models property.",
+                UserWarning,
+                stacklevel=2,
+            )
+            self._model = self._backend.model_name = name
+        else:
+            error_message = "No locally hosted model was found."
+            raise ValueError(error_message)
 
     def warm_up(self):
         """
@@ -91,6 +114,9 @@ def warm_up(self):
             model_kwargs=self._model_arguments,
         )
 
+        if not self.is_hosted and not self._model:
+            self.default_model()
+
     def to_dict(self) -> Dict[str, Any]:
         """
         Serializes the component to a dictionary.

diff --git a/integrations/nvidia/src/haystack_integrations/utils/nvidia/__init__.py b/integrations/nvidia/src/haystack_integrations/utils/nvidia/__init__.py
@@ -1,3 +1,4 @@
-from .utils import url_validation
+from .nim_backend import Model, NimBackend
+from .utils import is_hosted, url_validation
 
-__all__ = ["url_validation"]
+__all__ = ["NimBackend", "Model", "is_hosted", "url_validation"]