fix: fix Optimum types + add py.typed (#2018)

* fix: fix Optimum types + add py.typed

* check types with 3.10

* with 3.13

Author: anakin87

.github/workflows/optimum.yml

@@ -49,11 +49,11 @@ 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
+        # we check types with python 3.13 because with 3.9, the installation of some type stubs fails
+        # due to incompatibilities
+        if: matrix.python-version == '3.13' && runner.os == 'Linux'
+        run: hatch run fmt-check && hatch run test:types
 
       - name: Generate docs
         if: matrix.python-version == '3.9' && runner.os == 'Linux'

integrations/optimum/pyproject.toml

@@ -79,18 +79,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.optimum {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", "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.mypy.overrides]]
+module = "optimum.*"
+ignore_missing_imports = true
 
 [tool.hatch.metadata]
 allow-direct-references = true
@@ -173,26 +172,9 @@ omit = ["*/tests/*", "*/__init__.py"]
 show_missing = true
 exclude_lines = ["no cov", "if __name__ == .__main__.:", "if TYPE_CHECKING:"]
 
-
-[[tool.mypy.overrides]]
-module = [
-  "haystack.*",
-  "haystack_integrations.*",
-  "pytest.*",
-  "numpy.*",
-  "optimum.*",
-  "torch.*",
-  "transformers.*",
-  "huggingface_hub.*",
-  "sentence_transformers.*",
-]
-ignore_missing_imports = true
-
 [tool.pytest.ini_options]
 addopts = ["--strict-markers", "-vv"]
 markers = [
   "integration: integration tests",
-  "unit: unit tests",
-  "embedders: embedders tests",
 ]
 log_cli = true

integrations/optimum/src/haystack_integrations/components/embedders/optimum/_backend.py

@@ -2,7 +2,7 @@
 import json
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union, overload
 
 import numpy as np
 import torch
@@ -111,7 +111,7 @@ def __init__(self, params: _EmbedderParams):
         self.params = params
         self.model = None
         self.tokenizer = None
-        self.pooling_layer = None
+        self.pooling_layer: Optional[SentenceTransformerPoolingLayer] = None
 
     def warm_up(self):
         assert self.params.model_kwargs
@@ -188,6 +188,12 @@ def pool_embeddings(self, model_output: torch.Tensor, attention_mask: torch.Tens
         pooled_outputs = self.pooling_layer.forward(features)
         return pooled_outputs["sentence_embedding"]
 
+    @overload
+    def embed_texts(self, texts_to_embed: str) -> List[float]: ...
+
+    @overload
+    def embed_texts(self, texts_to_embed: List[str]) -> List[List[float]]: ...
+
     def embed_texts(
         self,
         texts_to_embed: Union[str, List[str]],

integrations/optimum/src/haystack_integrations/components/embedders/optimum/optimum_document_embedder.py

@@ -191,7 +191,7 @@ def _prepare_texts_to_embed(self, documents: List[Document]) -> List[str]:
         return texts_to_embed
 
     @component.output_types(documents=List[Document])
-    def run(self, documents: List[Document]):
+    def run(self, documents: List[Document]) -> Dict[str, List[Document]]:
         """
         Embed a list of Documents.
         The embedding of each Document is stored in the `embedding` field of the Document.

integrations/optimum/src/haystack_integrations/components/embedders/optimum/optimum_text_embedder.py

@@ -154,7 +154,7 @@ def from_dict(cls, data: Dict[str, Any]) -> "OptimumTextEmbedder":
         return default_from_dict(cls, data)
 
     @component.output_types(embedding=List[float])
-    def run(self, text: str):
+    def run(self, text: str) -> Dict[str, List[float]]:
         """
         Embed a string.