fix: mypy errors

libs/aws/langchain_aws/document_compressors/rerank.py

@@ -32,23 +32,21 @@ class BedrockRerank(BaseDocumentCompressor):
         arbitrary_types_allowed=True,
     )
 
-    @model_validator(mode="after")
-    def initialize_client(self) -> Self:
+    @model_validator(mode="before")
+    @classmethod
+    def initialize_client(cls, values: Dict[str, Any]) -> Any:
         """Initialize the AWS Bedrock client."""
-        if not self.client:
-            session = self._get_session()
-            self.client = session.client(
+        if not values.get("client"):
+            session = (
+                boto3.Session(profile_name=values.get("credentials_profile_name"))
+                if values.get("credentials_profile_name", None)
+                else boto3.Session()
+            )
+            values["client"] = session.client(
                 "bedrock-agent-runtime",
-                region_name=self.region_name
-                )
-        return self
-
-    def _get_session(self):
-        return (
-            boto3.Session(profile_name=self.credentials_profile_name)
-            if self.credentials_profile_name
-            else boto3.Session()
-        )
+                region_name=values.get("region_name"),
+            )
+        return values
 
     def rerank(
         self,

libs/aws/tests/unit_tests/document_compressors/test_rerank.py

@@ -7,7 +7,7 @@
 
 
 @pytest.fixture
-def reranker():
+def reranker() -> BedrockRerank:
     reranker = BedrockRerank(
         model_arn="arn:aws:bedrock:us-west-2::foundation-model/amazon.rerank-v1:0",
         region_name="us-east-1",
@@ -16,15 +16,14 @@ def reranker():
     return reranker
 
 @patch("boto3.Session")
-def test_initialize_client(mock_boto_session, reranker):
+def test_initialize_client(mock_boto_session: MagicMock, reranker: BedrockRerank) -> None:
     session_instance = MagicMock()
     mock_boto_session.return_value = session_instance
     session_instance.client.return_value = MagicMock()
-    reranker.initialize_client()
     assert reranker.client is not None
 
 @patch("langchain_aws.document_compressors.rerank.BedrockRerank.rerank")
-def test_rerank(mock_rerank, reranker):
+def test_rerank(mock_rerank: MagicMock, reranker: BedrockRerank) -> None:
     mock_rerank.return_value = [
         {"index": 0, "relevance_score": 0.9},
         {"index": 1, "relevance_score": 0.8},
@@ -41,7 +40,7 @@ def test_rerank(mock_rerank, reranker):
     assert results[1]["relevance_score"] == 0.8
 
 @patch("langchain_aws.document_compressors.rerank.BedrockRerank.rerank")
-def test_compress_documents(mock_rerank, reranker):
+def test_compress_documents(mock_rerank: MagicMock, reranker: BedrockRerank) -> None:
     mock_rerank.return_value = [
         {"index": 0, "relevance_score": 0.95},
         {"index": 1, "relevance_score": 0.85},