feat: Added setup logger

ai21/__init__.py

@@ -1,6 +1,7 @@
 from typing import Any
 
 from ai21.clients.studio.ai21_client import AI21Client
+from ai21.logger import setup_logger
 from ai21.resources.responses.answer_response import AnswerResponse
 from ai21.resources.responses.chat_response import ChatResponse
 from ai21.resources.responses.completion_response import CompletionsResponse
@@ -20,6 +21,7 @@
 from ai21.version import VERSION
 
 __version__ = VERSION
+setup_logger()
 
 
 def _import_bedrock_client():

ai21/ai21_env_config.py

@@ -17,6 +17,7 @@ class _AI21EnvConfig:
     timeout_sec: Optional[int] = None
     num_retries: Optional[int] = None
     aws_region: Optional[str] = None
+    log_level: Optional[str] = None
 
     @classmethod
     def from_env(cls) -> _AI21EnvConfig:
@@ -30,6 +31,7 @@ def from_env(cls) -> _AI21EnvConfig:
             timeout_sec=os.getenv("AI21_TIMEOUT_SEC"),
             num_retries=os.getenv("AI21_NUM_RETRIES"),
             aws_region=os.getenv("AI21_AWS_REGION", "us-east-1"),
+            log_level=os.getenv("AI21_LOG_LEVEL", "info"),
         )
 
 

ai21/clients/studio/resources/studio_answer.py

@@ -1,5 +1,6 @@
 from typing import Optional
 
+from ai21.logger import logger
 from ai21.resources.bases.answer_base import Answer
 from ai21.resources.responses.answer_response import AnswerResponse
 from ai21.resources.studio_resource import StudioResource

ai21/logger.py

@@ -1,3 +1,21 @@
 import logging
 
+from ai21.ai21_env_config import AI21EnvConfig
+
 logger = logging.getLogger("ai21")
+
+
+def _basic_config() -> None:
+    logging.basicConfig(
+        format="[%(asctime)s - %(name)s - %(levelname)s] %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
+
+
+def setup_logger() -> None:
+    _basic_config()
+
+    if AI21EnvConfig.log_level == "debug":
+        logger.setLevel(logging.DEBUG)
+    elif AI21EnvConfig.log_level == "info":
+        logger.setLevel(logging.INFO)