fix: fix Sagemaker types + add py.typed (#2023)

* fix: fix Sagemaker types + add py.typed

* improve type hint

Author: anakin87

.github/workflows/amazon_sagemaker.yml

@@ -49,11 +49,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/amazon_sagemaker/pyproject.toml

@@ -67,17 +67,20 @@ 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.generators.amazon_sagemaker {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.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.mypy.overrides]]
+module = [
+"botocore.*",
+"boto3.*",
+]
+ignore_missing_imports = true
 
 [tool.ruff]
 target-version = "py38"
@@ -154,24 +157,10 @@ omit = ["*/tests/*", "*/__init__.py"]
 show_missing = true
 exclude_lines = ["no cov", "if __name__ == .__main__.:", "if TYPE_CHECKING:"]
 
-[[tool.mypy.overrides]]
-module = [
-  "botocore.*",
-  "boto3.*",
-  "haystack.*",
-  "haystack_integrations.*",
-  "pytest.*",
-  "numpy.*",
-]
-ignore_missing_imports = true
-
 
 [tool.pytest.ini_options]
 addopts = "--strict-markers"
 markers = [
-  "unit: unit tests",
   "integration: integration tests",
-  "embedders: embedders tests",
-  "generators: generators tests",
 ]
 log_cli = true

integrations/amazon_sagemaker/src/haystack_integrations/components/generators/amazon_sagemaker/sagemaker.py

@@ -1,5 +1,5 @@
 import json
-from typing import Any, ClassVar, Dict, List, Optional
+from typing import Any, ClassVar, Dict, List, Optional, Union
 
 import boto3
 import requests
@@ -142,7 +142,7 @@ def to_dict(self) -> Dict[str, Any]:
         )
 
     @classmethod
-    def from_dict(cls, data) -> "SagemakerGenerator":
+    def from_dict(cls, data: Dict[str, Any]) -> "SagemakerGenerator":
         """
         Deserializes the component from a dictionary.
 
@@ -164,7 +164,7 @@ def _get_aws_session(
         aws_session_token: Optional[str] = None,
         aws_region_name: Optional[str] = None,
         aws_profile_name: Optional[str] = None,
-    ):
+    ) -> boto3.Session:
         """
         Creates an AWS Session with the given parameters.
 
@@ -192,7 +192,9 @@ def _get_aws_session(
             raise AWSConfigurationError(msg) from e
 
     @component.output_types(replies=List[str], meta=List[Dict[str, Any]])
-    def run(self, prompt: str, generation_kwargs: Optional[Dict[str, Any]] = None):
+    def run(
+        self, prompt: str, generation_kwargs: Optional[Dict[str, Any]] = None
+    ) -> Dict[str, Union[List[str], List[Dict[str, Any]]]]:
         """
         Invoke the text generation inference based on the provided prompt and generation parameters.