fix: Added static type checker

AI21Labs · Dec 14, 2023 · 3a48e8c · 3a48e8c
1 parent f00fd8c
commit 3a48e8c
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/ai21/clients/bedrock/ai21_bedrock_client.py b/ai21/clients/bedrock/ai21_bedrock_client.py
@@ -1,7 +1,6 @@
 import json
 import re
-from typing import Optional, Any, Dict
-
+from typing import Optional, Any, Dict, TYPE_CHECKING
 
 from ai21.ai21_env_config import AI21EnvConfig, _AI21EnvConfig
 from ai21.clients.bedrock.resources.bedrock_completion import BedrockCompletion
@@ -14,6 +13,9 @@
     r"Received client error \((.*?)\) from primary with message \"(.*?)\". "
     r"See .* in account .* for more information."
 )
+if TYPE_CHECKING:
+    import boto3
+    from botocore.exceptions import ClientError
 
 
 class AI21BedrockClient:

diff --git a/ai21/clients/sagemaker/ai21_sagemaker_client.py b/ai21/clients/sagemaker/ai21_sagemaker_client.py
@@ -1,6 +1,6 @@
 import json
 import re
-from typing import Optional
+from typing import Optional, TYPE_CHECKING
 
 from ai21.ai21_env_config import _AI21EnvConfig, AI21EnvConfig
 from ai21.clients.sagemaker.resources.sagemaker_answer import SageMakerAnswer
@@ -19,6 +19,10 @@
 )
 _SAGEMAKER_RUNTIME_NAME = "sagemaker-runtime"
 
+if TYPE_CHECKING:
+    import boto3
+    from botocore.exceptions import ClientError
+
 
 class AI21SageMakerClient:
     """