fix: Weaviate - fix types + add py.typed (#1977)

Author: anakin87

.github/workflows/weaviate.yml

@@ -44,11 +44,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: Run Weaviate container
         run: docker compose up -d

integrations/weaviate/pyproject.toml

@@ -70,18 +70,14 @@ 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}"
-
-# 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.hatch.envs.lint.scripts]
-typing = "mypy --install-types --non-interactive --explicit-package-bases {args:src/ tests}"
+types = """mypy -p haystack_integrations.document_stores.weaviate \
+-p haystack_integrations.components.retrievers.weaviate {args}"""
+
+[tool.mypy]
+install_types = true
+non_interactive = true
+check_untyped_defs = true
+disallow_incomplete_defs = true
 
 [tool.black]
 target-version = ["py38"]

integrations/weaviate/src/haystack_integrations/components/retrievers/py.typed

@@ -89,7 +89,9 @@ def from_dict(cls, data: Dict[str, Any]) -> "WeaviateBM25Retriever":
         return default_from_dict(cls, data)
 
     @component.output_types(documents=List[Document])
-    def run(self, query: str, filters: Optional[Dict[str, Any]] = None, top_k: Optional[int] = None):
+    def run(
+        self, query: str, filters: Optional[Dict[str, Any]] = None, top_k: Optional[int] = None
+    ) -> Dict[str, List[Document]]:
         """
         Retrieves documents from Weaviate using the BM25 algorithm.
 

integrations/weaviate/src/haystack_integrations/components/retrievers/weaviate/bm25_retriever.py

@@ -106,7 +106,7 @@ def run(
         top_k: Optional[int] = None,
         distance: Optional[float] = None,
         certainty: Optional[float] = None,
-    ):
+    ) -> Dict[str, List[Document]]:
         """
         Retrieves documents from Weaviate using the vector search.
 

integrations/weaviate/src/haystack_integrations/components/retrievers/weaviate/embedding_retriever.py

@@ -30,7 +30,7 @@ def __str__(self):
         return self.value
 
     @staticmethod
-    def from_class(auth_class) -> "SupportedAuthTypes":
+    def from_class(auth_class: Type["AuthCredentials"]) -> "SupportedAuthTypes":
         auth_types = {
             AuthApiKey: SupportedAuthTypes.API_KEY,
             AuthBearerToken: SupportedAuthTypes.BEARER,
@@ -80,7 +80,7 @@ def from_dict(data: Dict[str, Any]) -> "AuthCredentials":
 
     @classmethod
     @abstractmethod
-    def _from_dict(cls, data: Dict[str, Any]):
+    def _from_dict(cls, data: Dict[str, Any]) -> "AuthCredentials":
         """
         Internal method to convert a dictionary representation to an auth credentials object.
         All subclasses must implement this method.
@@ -109,7 +109,8 @@ def _from_dict(cls, data: Dict[str, Any]) -> "AuthApiKey":
         return cls(**data["init_parameters"])
 
     def resolve_value(self) -> WeaviateAuthApiKey:
-        return WeaviateAuthApiKey(api_key=self.api_key.resolve_value())
+        # resolve_value, when used with Secret.from_env_var (strict=True), returns a string or raises an error
+        return WeaviateAuthApiKey(api_key=self.api_key.resolve_value())  # type: ignore[arg-type]
 
 
 @dataclass(frozen=True)
@@ -136,9 +137,11 @@ def resolve_value(self) -> WeaviateAuthBearerToken:
         refresh_token = self.refresh_token.resolve_value()
 
         return WeaviateAuthBearerToken(
-            access_token=access_token,
+            # resolve_value, when used with Secret.from_env_var (strict=True), returns a string or raises an error
+            access_token=access_token,  # type: ignore[arg-type]
             expires_in=self.expires_in,
-            refresh_token=refresh_token,
+            # resolve_value, when used with Secret.from_env_var (strict=False), returns a string or None
+            refresh_token=refresh_token,  # type: ignore[arg-type]
         )
 
 
@@ -162,8 +165,10 @@ def _from_dict(cls, data: Dict[str, Any]) -> "AuthClientCredentials":
 
     def resolve_value(self) -> WeaviateAuthClientCredentials:
         return WeaviateAuthClientCredentials(
-            client_secret=self.client_secret.resolve_value(),
-            scope=self.scope.resolve_value(),
+            # resolve_value, when used with Secret.from_env_var (strict=True), returns a string or raises an error
+            client_secret=self.client_secret.resolve_value(),  # type: ignore[arg-type]
+            # resolve_value, when used with Secret.from_env_var (strict=False), returns a string or None
+            scope=self.scope.resolve_value(),  # type: ignore[arg-type]
         )
 
 
@@ -189,7 +194,9 @@ def _from_dict(cls, data: Dict[str, Any]) -> "AuthClientPassword":
 
     def resolve_value(self) -> WeaviateAuthClientPassword:
         return WeaviateAuthClientPassword(
-            username=self.username.resolve_value(),
-            password=self.password.resolve_value(),
-            scope=self.scope.resolve_value(),
+            # resolve_value, when used with Secret.from_env_var (strict=True), returns a string or raises an error
+            username=self.username.resolve_value(),  # type: ignore[arg-type]
+            password=self.password.resolve_value(),  # type: ignore[arg-type]
+            # resolve_value, when used with Secret.from_env_var (strict=False), returns a string or None
+            scope=self.scope.resolve_value(),  # type: ignore[arg-type]
         )

integrations/weaviate/src/haystack_integrations/document_stores/py.typed

@@ -144,8 +144,8 @@ def __init__(
         self._additional_config = additional_config
         self._grpc_port = grpc_port
         self._grpc_secure = grpc_secure
-        self._client = None
-        self._collection = None
+        self._client: Optional[weaviate.WeaviateClient] = None
+        self._collection: Optional[weaviate.Collection] = None
         # Store the connection settings dictionary
         self._collection_settings = collection_settings or {
             "class": "Default",
@@ -173,9 +173,12 @@ def client(self):
             # If we detect that the URL is a Weaviate Cloud URL, we use the utility function to connect
             # instead of using WeaviateClient directly like in other cases.
             # Among other things, the utility function takes care of parsing the URL.
+            if not self._auth_client_secret:
+                msg = "Auth credentials are required for Weaviate Cloud Services"
+                raise ValueError(msg)
             self._client = weaviate.connect_to_weaviate_cloud(
                 self._url,
-                auth_credentials=self._auth_client_secret.resolve_value() if self._auth_client_secret else None,
+                auth_credentials=self._auth_client_secret.resolve_value(),
                 headers=self._additional_headers,
                 additional_config=self._additional_config,
             )
@@ -343,7 +346,7 @@ def _to_document(self, data: DataObject[Dict[str, Any], None]) -> Document:
 
         return Document.from_dict(document_data)
 
-    def _query(self) -> List[Dict[str, Any]]:
+    def _query(self) -> List[DataObject[Dict[str, Any], None]]:
         properties = [p.name for p in self.collection.config.get().properties]
         try:
             result = self.collection.iterator(include_vector=True, return_properties=properties)
@@ -352,7 +355,7 @@ def _query(self) -> List[Dict[str, Any]]:
             raise DocumentStoreError(msg) from e
         return result
 
-    def _query_with_filters(self, filters: Dict[str, Any]) -> List[Dict[str, Any]]:
+    def _query_with_filters(self, filters: Dict[str, Any]) -> List[DataObject[Dict[str, Any], None]]:
         properties = [p.name for p in self.collection.config.get().properties]
         # When querying with filters we need to paginate using limit and offset as using
         # a cursor with after is not possible. See the official docs: