Add default model for NVIDIA HayStack local NIM endpoints

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", "pydantic"]
 [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
 

integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/_nim_backend.py

@@ -3,7 +3,7 @@
 import requests
 from haystack.utils import Secret
 
-from .backend import EmbedderBackend
+from .backend import EmbedderBackend, Model
 
 REQUEST_TIMEOUT = 60
 
@@ -50,3 +50,20 @@ def embed(self, texts: List[str]) -> Tuple[List[List[float]], Dict[str, Any]]:
         embeddings = [e["embedding"] for e in sorted(data["data"], key=lambda e: e["index"])]
 
         return embeddings, {"usage": data["usage"]}
+
+    def models(self) -> List[Model]:
+        url = f"{self.api_url}/models"
+
+        res = self.session.get(
+            url,
+            timeout=REQUEST_TIMEOUT,
+        )
+        res.raise_for_status()
+
+        data = res.json()["data"]
+        models = []
+        for element in data:
+            assert "id" in element, f"No id found in {element}"
+            models.append(Model(id=element["id"]))
+
+        return models

integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/backend.py

@@ -1,6 +1,24 @@
 from abc import ABC, abstractmethod
 from typing import Any, Dict, List, Optional, Tuple
 
+from pydantic import BaseModel, ConfigDict
+
+
+class Model(BaseModel):
+    """
+    Model information.
+
+    id: unique identifier for the model, passed as model parameter for requests
+    aliases: list of aliases for the model
+    base_model: root model for the model
+    All aliases are deprecated and will trigger a warning when used.
+    """
+
+    id: str
+    model_config = ConfigDict(from_attributes=True, protected_namespaces=())
+    aliases: Optional[list] = None
+    base_model: Optional[str] = None
+
 
 class EmbedderBackend(ABC):
     def __init__(self, model: str, model_kwargs: Optional[Dict[str, Any]] = None):
@@ -27,3 +45,13 @@ def embed(self, texts: List[str]) -> Tuple[List[List[float]], Dict[str, Any]]:
             metadata returned by the service.
         """
         pass
+
+    @abstractmethod
+    def models(self) -> List[Model]:
+        """
+        Invoke the backend to get available models.
+
+        :return:
+            Models available
+        """
+        pass

integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/document_embedder.py

@@ -1,4 +1,6 @@
+import warnings
 from typing import Any, Dict, List, Optional, Tuple, Union
+from urllib.parse import urlparse
 
 from haystack import Document, component, default_from_dict, default_to_dict
 from haystack.utils import Secret, deserialize_secrets_inplace
@@ -31,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 = "https://ai.api.nvidia.com/v1/retrieval/nvidia",
         prefix: str = "",
@@ -85,6 +87,32 @@ def __init__(
 
         self.backend: Optional[EmbedderBackend] = None
         self._initialized = False
+        self.is_hosted = urlparse(self.api_url).netloc in [
+            "integrate.api.nvidia.com",
+            "ai.api.nvidia.com",
+        ]
+        if self.is_hosted 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):
         """
@@ -104,6 +132,8 @@ def warm_up(self):
         )
 
         self._initialized = True
+        if not self.is_hosted and not self.model:
+            self.default_model()
 
     def to_dict(self) -> Dict[str, Any]:
         """

integrations/nvidia/src/haystack_integrations/components/embedders/nvidia/text_embedder.py

@@ -1,4 +1,6 @@
+import warnings
 from typing import Any, Dict, List, Optional, Union
+from urllib.parse import urlparse
 
 from haystack import component, default_from_dict, default_to_dict
 from haystack.utils import Secret, deserialize_secrets_inplace
@@ -32,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 = "https://ai.api.nvidia.com/v1/retrieval/nvidia",
         prefix: str = "",
@@ -70,6 +72,33 @@ def __init__(
         self.backend: Optional[EmbedderBackend] = None
         self._initialized = False
 
+        self.is_hosted = urlparse(self.api_url).netloc in [
+            "integrate.api.nvidia.com",
+            "ai.api.nvidia.com",
+        ]
+        if self.is_hosted 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.
@@ -89,6 +118,9 @@ def warm_up(self):
 
         self._initialized = True
 
+        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.

integrations/nvidia/src/haystack_integrations/components/generators/nvidia/_nim_backend.py

@@ -3,7 +3,7 @@
 import requests
 from haystack.utils import Secret
 
-from .backend import GeneratorBackend
+from .backend import GeneratorBackend, Model
 
 REQUEST_TIMEOUT = 60
 
@@ -78,3 +78,20 @@ def generate(self, prompt: str) -> Tuple[List[str], List[Dict[str, Any]]]:
             meta.append(choice_meta)
 
         return replies, meta
+
+    def models(self) -> List[Model]:
+        url = f"{self.api_url}/models"
+
+        res = self.session.get(
+            url,
+            timeout=REQUEST_TIMEOUT,
+        )
+        res.raise_for_status()
+
+        data = res.json()["data"]
+        models = []
+        for element in data:
+            assert "id" in element, f"No id found in {element}"
+            models.append(Model(id=element["id"]))
+
+        return models

integrations/nvidia/src/haystack_integrations/components/generators/nvidia/backend.py

@@ -1,6 +1,24 @@
 from abc import ABC, abstractmethod
 from typing import Any, Dict, List, Optional, Tuple
 
+from pydantic import BaseModel, ConfigDict
+
+
+class Model(BaseModel):
+    """
+    Model information.
+
+    id: unique identifier for the model, passed as model parameter for requests
+    aliases: list of aliases for the model
+    base_model: root model for the model
+    All aliases are deprecated and will trigger a warning when used.
+    """
+
+    id: str
+    model_config = ConfigDict(from_attributes=True, protected_namespaces=())
+    aliases: Optional[list] = None
+    base_model: Optional[str] = None
+
 
 class GeneratorBackend(ABC):
     def __init__(self, model: str, model_kwargs: Optional[Dict[str, Any]] = None):
@@ -27,3 +45,13 @@ def generate(self, prompt: str) -> Tuple[List[str], List[Dict[str, Any]]]:
             metadata returned by the service.
         """
         pass
+
+    @abstractmethod
+    def models(self) -> List[Model]:
+        """
+        Invoke the backend to get available models.
+
+        :return:
+            Models available
+        """
+        pass

integrations/nvidia/src/haystack_integrations/components/generators/nvidia/generator.py

@@ -1,7 +1,9 @@
 # SPDX-FileCopyrightText: 2024-present deepset GmbH <[email protected]>
 #
 # SPDX-License-Identifier: Apache-2.0
+import warnings
 from typing import Any, Dict, List, Optional
+from urllib.parse import urlparse
 
 from haystack import component, default_from_dict, default_to_dict
 from haystack.utils.auth import Secret, deserialize_secrets_inplace
@@ -41,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,
@@ -69,6 +71,30 @@ def __init__(
 
         self._backend: Optional[GeneratorBackend] = None
 
+        self.is_hosted = urlparse(api_url).netloc in [
+            "integrate.api.nvidia.com",
+            "ai.api.nvidia.com",
+        ]
+
+    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):
         """
         Initializes the component.
@@ -86,6 +112,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.

integrations/nvidia/tests/conftest.py

@@ -0,0 +1,20 @@
+import pytest
+from requests_mock import Mocker
+
+
+@pytest.fixture
+def mock_local_models(requests_mock: Mocker) -> None:
+    requests_mock.get(
+        "http://localhost:8080/v1/models",
+        json={
+            "data": [
+                {
+                    "id": "model1",
+                    "object": "model",
+                    "created": 1234567890,
+                    "owned_by": "OWNER",
+                    "root": "model1",
+                },
+            ]
+        },
+    )