From 3a48e8c0f4540949b9a99ba031a6e82347a7c91a Mon Sep 17 00:00:00 2001
From: Asaf Gardin <asafg@ai21.com>
Date: Thu, 14 Dec 2023 13:34:17 +0200
Subject: [PATCH] fix: Added static type checker

---
 ai21/clients/bedrock/ai21_bedrock_client.py     | 6 ++++--
 ai21/clients/sagemaker/ai21_sagemaker_client.py | 6 +++++-
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/ai21/clients/bedrock/ai21_bedrock_client.py b/ai21/clients/bedrock/ai21_bedrock_client.py
index 985cce28..4dfb39a4 100644
--- 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
index b76fdd64..bcddc373 100644
--- 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:
     """