fix: fix Fastembed types + add py.typed (#1954)

Author: anakin87

.github/workflows/fastembed.yml

@@ -33,12 +33,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/fastembed/pyproject.toml

@@ -65,17 +65,14 @@ unit = 'pytest -m "not integration" {args:tests}'
 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.fastembed \
+-p haystack_integrations.components.rankers.fastembed {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", "numpy"]
-[tool.hatch.envs.lint.scripts]
-typing = "mypy --install-types --non-interactive --explicit-package-bases {args:src/ tests}"
+[tool.mypy]
+install_types = true
+non_interactive = true
+check_untyped_defs = true
+disallow_incomplete_defs = true
 
 [tool.black]
 target-version = ["py38"]
@@ -161,14 +158,5 @@ exclude_lines = ["no cov", "if __name__ == .__main__.:", "if TYPE_CHECKING:"]
 
 [tool.pytest.ini_options]
 minversion = "6.0"
-markers = ["unit: unit tests", "integration: integration tests"]
-
-[[tool.mypy.overrides]]
-module = [
-  "haystack.*",
-  "haystack_integrations.*",
-  "fastembed.*",
-  "pytest.*",
-  "numpy.*",
-]
-ignore_missing_imports = true
+markers = ["integration: integration tests"]
+

integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/embedding_backend/fastembed_backend.py

@@ -20,7 +20,7 @@ def get_embedding_backend(
         cache_dir: Optional[str] = None,
         threads: Optional[int] = None,
         local_files_only: bool = False,
-    ):
+    ) -> "_FastembedEmbeddingBackend":
         embedding_backend_id = f"{model_name}{cache_dir}{threads}"
 
         if embedding_backend_id in _FastembedEmbeddingBackendFactory._instances:
@@ -49,7 +49,7 @@ def __init__(
             model_name=model_name, cache_dir=cache_dir, threads=threads, local_files_only=local_files_only
         )
 
-    def embed(self, data: List[str], progress_bar=True, **kwargs) -> List[List[float]]:
+    def embed(self, data: List[str], progress_bar: bool = True, **kwargs: Any) -> List[List[float]]:
         # the embed method returns a Iterable[np.ndarray], so we convert it to a list of lists
         embeddings = []
         embeddings_iterable = self.model.embed(data, **kwargs)
@@ -74,7 +74,7 @@ def get_embedding_backend(
         threads: Optional[int] = None,
         local_files_only: bool = False,
         model_kwargs: Optional[Dict[str, Any]] = None,
-    ):
+    ) -> "_FastembedSparseEmbeddingBackend":
         embedding_backend_id = f"{model_name}{cache_dir}{threads}{local_files_only}{model_kwargs}"
 
         if embedding_backend_id in _FastembedSparseEmbeddingBackendFactory._instances:
@@ -114,7 +114,7 @@ def __init__(
             **model_kwargs,
         )
 
-    def embed(self, data: List[List[str]], progress_bar=True, **kwargs) -> List[SparseEmbedding]:
+    def embed(self, data: List[str], progress_bar: bool = True, **kwargs: Any) -> List[SparseEmbedding]:
         # The embed method returns a Iterable[SparseEmbedding], so we convert to Haystack SparseEmbedding type.
         # Each SparseEmbedding contains an `indices` key containing a list of int and
         # an `values` key containing a list of floats.

integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_document_embedder.py

@@ -150,7 +150,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]]:
         """
         Embeds a list of Documents.
 

integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_sparse_document_embedder.py

@@ -141,7 +141,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]]:
         """
         Embeds a list of Documents.
 

integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_sparse_text_embedder.py

@@ -94,7 +94,7 @@ def warm_up(self):
             )
 
     @component.output_types(sparse_embedding=SparseEmbedding)
-    def run(self, text: str):
+    def run(self, text: str) -> Dict[str, SparseEmbedding]:
         """
         Embeds text using the Fastembed model.
 

integrations/fastembed/src/haystack_integrations/components/embedders/fastembed/fastembed_text_embedder.py

@@ -96,7 +96,7 @@ def warm_up(self):
             )
 
     @component.output_types(embedding=List[float])
-    def run(self, text: str):
+    def run(self, text: str) -> Dict[str, List[float]]:
         """
         Embeds text using the Fastembed model.
 

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

@@ -76,7 +76,7 @@ def __init__(
         self.local_files_only = local_files_only
         self.meta_fields_to_embed = meta_fields_to_embed or []
         self.meta_data_separator = meta_data_separator
-        self._model = None
+        self._model: Optional[TextCrossEncoder] = None
 
     def to_dict(self) -> Dict[str, Any]:
         """
@@ -140,7 +140,7 @@ def _prepare_fastembed_input_docs(self, documents: List[Document]) -> List[str]:
         return concatenated_input_list
 
     @component.output_types(documents=List[Document])
-    def run(self, query: str, documents: List[Document], top_k: Optional[int] = None):
+    def run(self, query: str, documents: List[Document], top_k: Optional[int] = None) -> Dict[str, List[Document]]:
         """
         Returns a list of documents ranked by their similarity to the given query, using FastEmbed.