fix: fix Nvidia types + add py.typed (#1970)

* draft

* client enum

* fix workflow

* improvements

* Apply suggestions from code review

Co-authored-by: Julian Risch <[email protected]>

* fixes

* model_type optional + NimBackend tests

* tests

---------

Co-authored-by: Julian Risch <[email protected]>

Author: anakin87

.github/workflows/nvidia.yml

@@ -51,11 +51,9 @@ jobs:
       - name: Install Hatch
         run: pip install --upgrade hatch
 
-    # TODO: Once this integration is properly typed, use hatch run test:types
-    # https://github.com/deepset-ai/haystack-core-integrations/issues/1771
       - name: Lint
         if: matrix.python-version == '3.9' && runner.os == 'Linux'
-        run: hatch run fmt-check && hatch run lint:typing
+        run: hatch run fmt-check && hatch run test:types
 
       - name: Generate docs
         if: matrix.python-version == '3.9' && runner.os == 'Linux'

integrations/nvidia/pyproject.toml

@@ -68,18 +68,17 @@ integration = 'pytest -m "integration" {args:tests}'
 all = 'pytest {args:tests}'
 cov-retry = 'all --cov=haystack_integrations --reruns 3 --reruns-delay 30 -x'
 
-types = "mypy --install-types --non-interactive --explicit-package-bases {args:src/ tests}"
+types = """mypy -p haystack_integrations.components.embedders.nvidia \
+-p haystack_integrations.components.generators.nvidia \
+-p haystack_integrations.components.rankers.nvidia \
+-p haystack_integrations.utils.nvidia {args}"""
 
-# TODO: remove lint environment once this integration is properly typed
-# test environment should be used instead
-# https://github.com/deepset-ai/haystack-core-integrations/issues/1771
-[tool.hatch.envs.lint]
-installer = "uv"
-detached = true
-dependencies = ["pip", "black>=23.1.0", "mypy>=1.0.0", "ruff>=0.0.243"]
+[tool.mypy]
+install_types = true
+non_interactive = true
+check_untyped_defs = true
+disallow_incomplete_defs = true
 
-[tool.hatch.envs.lint.scripts]
-typing = "mypy --install-types --non-interactive --explicit-package-bases {args:src/ tests}"
 
 [tool.black]
 target-version = ["py38"]
@@ -160,27 +159,9 @@ omit = ["*/tests/*", "*/__init__.py"]
 show_missing = true
 exclude_lines = ["no cov", "if __name__ == .__main__.:", "if TYPE_CHECKING:"]
 
-
-[[tool.mypy.overrides]]
-module = [
-  "nvidia.*",
-  "haystack.*",
-  "haystack_integrations.*",
-  "pytest.*",
-  "numpy.*",
-  "requests_mock.*",
-  "openai.*",
-  "pydantic.*",
-]
-ignore_missing_imports = true
-
 [tool.pytest.ini_options]
 addopts = "--strict-markers"
 markers = [
   "integration: integration tests",
-  "unit: unit tests",
-  "embedders: embedders tests",
-  "generators: generators tests",
-  "chat_generators: chat_generators tests",
 ]
 log_cli = true

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

@@ -11,7 +11,7 @@
 from tqdm import tqdm
 
 from haystack_integrations.components.embedders.nvidia.truncate import EmbeddingTruncateMode
-from haystack_integrations.utils.nvidia import DEFAULT_API_URL, Model, NimBackend, url_validation
+from haystack_integrations.utils.nvidia import DEFAULT_API_URL, Client, Model, NimBackend, url_validation
 
 logger = logging.getLogger(__name__)
 
@@ -122,7 +122,9 @@ def default_model(self):
                 UserWarning,
                 stacklevel=2,
             )
-            self.model = self.backend.model = name
+            self.model = name
+            if self.backend:
+                self.backend.model = name
         else:
             error_message = "No locally hosted model was found."
             raise ValueError(error_message)
@@ -143,7 +145,7 @@ def warm_up(self):
             api_url=self.api_url,
             api_key=self.api_key,
             model_kwargs=model_kwargs,
-            client=self.__class__.__name__,
+            client=Client.NVIDIA_DOCUMENT_EMBEDDER,
             timeout=self.timeout,
         )
         if not self.model and self.backend.model:
@@ -232,7 +234,7 @@ def _embed_batch(self, texts_to_embed: List[str], batch_size: int) -> Tuple[List
         return all_embeddings, {"usage": {"prompt_tokens": usage_prompt_tokens, "total_tokens": usage_total_tokens}}
 
     @component.output_types(documents=List[Document], meta=Dict[str, Any])
-    def run(self, documents: List[Document]):
+    def run(self, documents: List[Document]) -> Dict[str, Union[List[Document], Dict[str, Any]]]:
         """
         Embed a list of Documents.
 

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

@@ -10,7 +10,7 @@
 from haystack.utils import Secret, deserialize_secrets_inplace
 
 from haystack_integrations.components.embedders.nvidia.truncate import EmbeddingTruncateMode
-from haystack_integrations.utils.nvidia import DEFAULT_API_URL, Model, NimBackend, url_validation
+from haystack_integrations.utils.nvidia import DEFAULT_API_URL, Client, Model, NimBackend, url_validation
 
 logger = logging.getLogger(__name__)
 
@@ -112,7 +112,9 @@ def default_model(self):
                 UserWarning,
                 stacklevel=2,
             )
-            self.model = self.backend.model = name
+            self.model = name
+            if self.backend:
+                self.backend.model = name
         else:
             error_message = "No locally hosted model was found."
             raise ValueError(error_message)
@@ -134,7 +136,7 @@ def warm_up(self):
             api_key=self.api_key,
             model_kwargs=model_kwargs,
             timeout=self.timeout,
-            client=self.__class__.__name__,
+            client=Client.NVIDIA_TEXT_EMBEDDER,
         )
         self._initialized = True
 
@@ -185,7 +187,7 @@ def from_dict(cls, data: Dict[str, Any]) -> "NvidiaTextEmbedder":
         return default_from_dict(cls, data)
 
     @component.output_types(embedding=List[float], meta=Dict[str, Any])
-    def run(self, text: str):
+    def run(self, text: str) -> Dict[str, Union[List[float], Dict[str, Any]]]:
         """
         Embed a string.
 

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

@@ -4,12 +4,12 @@
 
 import os
 import warnings
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 
 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 DEFAULT_API_URL, Model, NimBackend, is_hosted, url_validation
+from haystack_integrations.utils.nvidia import DEFAULT_API_URL, Client, Model, NimBackend, is_hosted, url_validation
 
 
 @component
@@ -104,7 +104,9 @@ def default_model(self):
                 UserWarning,
                 stacklevel=2,
             )
-            self._model = self.backend.model = name
+            self._model = name
+            if self.backend:
+                self.backend.model = name
         else:
             error_message = "No locally hosted model was found."
             raise ValueError(error_message)
@@ -123,7 +125,7 @@ def warm_up(self):
             api_key=self._api_key,
             model_kwargs=self._model_arguments,
             timeout=self.timeout,
-            client=self.__class__.__name__,
+            client=Client.NVIDIA_GENERATOR,
         )
 
         if not self.is_hosted and not self._model:
@@ -169,7 +171,7 @@ def from_dict(cls, data: Dict[str, Any]) -> "NvidiaGenerator":
         return default_from_dict(cls, data)
 
     @component.output_types(replies=List[str], meta=List[Dict[str, Any]])
-    def run(self, prompt: str):
+    def run(self, prompt: str) -> Dict[str, Union[List[str], List[Dict[str, Any]]]]:
         """
         Queries the model with the provided prompt.
 

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

@@ -10,7 +10,7 @@
 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, NimBackend, is_hosted, url_validation
+from haystack_integrations.utils.nvidia import DEFAULT_API_URL, Client, NimBackend, is_hosted, url_validation
 
 logger = logging.getLogger(__name__)
 
@@ -162,7 +162,7 @@ def warm_up(self):
         :raises ValueError: If the API key is required for hosted NVIDIA NIMs.
         """
         if not self._initialized:
-            model_kwargs = {}
+            model_kwargs: Dict[str, Any] = {}
             if self.truncate is not None:
                 model_kwargs.update(truncate=str(self.truncate))
             self.backend = NimBackend(
@@ -172,9 +172,9 @@ def warm_up(self):
                 api_key=self.api_key,
                 model_kwargs=model_kwargs,
                 timeout=self.timeout,
-                client=self.__class__.__name__,
+                client=Client.NVIDIA_RANKER,
             )
-            if not self.is_hosted and not self._model:
+            if not self.is_hosted and not self.model:
                 if self.backend.model:
                     self.model = self.backend.model
             self._initialized = True

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

@@ -2,8 +2,9 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+from .client import Client
 from .models import DEFAULT_API_URL, Model
 from .nim_backend import NimBackend
 from .utils import is_hosted, url_validation
 
-__all__ = ["DEFAULT_API_URL", "Model", "NimBackend", "is_hosted", "url_validation", "validate_hosted_model"]
+__all__ = ["DEFAULT_API_URL", "Client", "Model", "NimBackend", "is_hosted", "url_validation", "validate_hosted_model"]

integrations/nvidia/src/haystack_integrations/components/rankers/nvidia/py.typed

@@ -0,0 +1,26 @@
+from enum import Enum
+
+
+class Client(Enum):
+    """
+    Client to use for NVIDIA NIMs.
+    """
+
+    NVIDIA_GENERATOR = "NvidiaGenerator"
+    NVIDIA_TEXT_EMBEDDER = "NvidiaTextEmbedder"
+    NVIDIA_DOCUMENT_EMBEDDER = "NvidiaDocumentEmbedder"
+    NVIDIA_RANKER = "NvidiaRanker"
+
+    def __str__(self) -> str:
+        """Convert a Client enum to a string."""
+        return self.value
+
+    @staticmethod
+    def from_str(string: str) -> "Client":
+        """Convert a string to a Client enum."""
+        enum_map = {e.value: e for e in Client}
+        mode = enum_map.get(string)
+        if mode is None:
+            msg = f"Unknown client '{string}' to use for NVIDIA NIMs. Supported modes are: {list(enum_map.keys())}"
+            raise ValueError(msg)
+        return mode