diff --git a/ai21/ai21_env_config.py b/ai21/ai21_env_config.py index fabcb34e..e6e1b738 100644 --- a/ai21/ai21_env_config.py +++ b/ai21/ai21_env_config.py @@ -1,4 +1,6 @@ from __future__ import annotations + +import logging import os from dataclasses import dataclass from typing import Optional @@ -14,6 +16,8 @@ _ENV_AWS_REGION = "AI21_AWS_REGION" _ENV_LOG_LEVEL = "AI21_LOG_LEVEL" +_logger = logging.getLogger(__name__) + @dataclass class _AI21EnvConfig: @@ -80,5 +84,20 @@ def log_level(self) -> Optional[str]: self._log_level = os.getenv(_ENV_LOG_LEVEL, self._log_level) return self._log_level + def log(self, with_secrets: bool = False) -> None: + env_vars = { + _ENV_API_VERSION: self.api_version, + _ENV_API_HOST: self.api_host, + _ENV_TIMEOUT_SEC: self.timeout_sec, + _ENV_NUM_RETRIES: self.num_retries, + _ENV_AWS_REGION: self.aws_region, + _ENV_LOG_LEVEL: self.log_level, + } + + if with_secrets: + env_vars[_ENV_API_KEY] = self.api_key + + _logger.debug(f"AI21 environment configuration: {env_vars}") + AI21EnvConfig = _AI21EnvConfig.from_env() diff --git a/ai21/http_client/async_http_client.py b/ai21/http_client/async_http_client.py index 914df7cd..eaeae50a 100644 --- a/ai21/http_client/async_http_client.py +++ b/ai21/http_client/async_http_client.py @@ -82,7 +82,7 @@ async def _request( ) -> httpx.Response: timeout = self._timeout_sec headers = self._headers - logger.debug(f"Calling {method} {url} {headers} {params}") + logger.debug(f"Calling {method} {url} {headers} {params} {body}") if method == "GET": request = self._client.build_request( diff --git a/ai21/http_client/base_http_client.py b/ai21/http_client/base_http_client.py index 32b62797..bf9efe98 100644 --- a/ai21/http_client/base_http_client.py +++ b/ai21/http_client/base_http_client.py @@ -1,14 +1,11 @@ from __future__ import annotations import json - -from typing import Generic, TypeVar, Union, Any, Optional, Dict, BinaryIO from abc import ABC, abstractmethod +from typing import Generic, TypeVar, Union, Any, Optional, Dict, BinaryIO import httpx -from ai21.stream.stream import Stream -from ai21.stream.async_stream import AsyncStream from ai21.errors import ( BadRequest, Unauthorized, @@ -18,11 +15,12 @@ ServiceUnavailable, AI21APIError, ) +from ai21.stream.async_stream import AsyncStream +from ai21.stream.stream import Stream _HttpxClientT = TypeVar("_HttpxClientT", bound=Union[httpx.Client, httpx.AsyncClient]) _DefaultStreamT = TypeVar("_DefaultStreamT", bound=Union[Stream[Any], AsyncStream[Any]]) - DEFAULT_TIMEOUT_SEC = 300 DEFAULT_NUM_RETRIES = 0 RETRY_BACK_OFF_FACTOR = 0.5 diff --git a/ai21/http_client/http_client.py b/ai21/http_client/http_client.py index b7500806..770bb15d 100644 --- a/ai21/http_client/http_client.py +++ b/ai21/http_client/http_client.py @@ -82,7 +82,7 @@ def _request( ) -> httpx.Response: timeout = self._timeout_sec headers = self._headers - logger.debug(f"Calling {method} {url} {headers} {params}") + logger.debug(f"Calling {method} {url} {headers} {params} {body}") if method == "GET": request = self._client.build_request( diff --git a/ai21/logger.py b/ai21/logger.py index 49621c03..a0706e5b 100644 --- a/ai21/logger.py +++ b/ai21/logger.py @@ -1,8 +1,64 @@ import logging +import os +import re from ai21.ai21_env_config import AI21EnvConfig +_verbose = False + logger = logging.getLogger("ai21") +httpx_logger = logging.getLogger("httpx") + + +class CensorSecretsFormatter(logging.Formatter): + def format(self, record: logging.LogRecord) -> str: + # Get original log message + message = super().format(record) + + if not get_verbose(): + return self._censor_secrets(message) + + return message + + def _censor_secrets(self, message: str) -> str: + # Regular expression to find the Authorization key and its value + pattern = r"('Authorization':\s*'[^']*'|'api-key':\s*'[^']*'|'X-Amz-Security-Token':\s*'[^']*')" + + def replacement(match): + return match.group(0).split(":")[0] + ": '**************'" + + # Substitute the Authorization value with ************** + return re.sub(pattern, replacement, message) + + +def set_verbose(value: bool) -> None: + """ + Use this function if you want to log additional, more sensitive data like - secrets and environment variables. + Log level will be set to DEBUG if verbose is set to True. + """ + global _verbose + _verbose = value + + set_debug(_verbose) + + AI21EnvConfig.log(with_secrets=value) + + +def set_debug(value: bool) -> None: + """ + Additional way to set log level to DEBUG. + """ + if value: + os.environ["AI21_LOG_LEVEL"] = "debug" + else: + os.environ["AI21_LOG_LEVEL"] = "info" + + setup_logger() + + +def get_verbose() -> bool: + global _verbose + return _verbose def _basic_config() -> None: @@ -14,8 +70,12 @@ def _basic_config() -> None: def setup_logger() -> None: _basic_config() + # Set the root handler with the censor formatter + logger.root.handlers[0].setFormatter(CensorSecretsFormatter()) - if AI21EnvConfig.log_level == "debug": + if AI21EnvConfig.log_level.lower() == "debug": logger.setLevel(logging.DEBUG) - elif AI21EnvConfig.log_level == "info": + httpx_logger.setLevel(logging.DEBUG) + elif AI21EnvConfig.log_level.lower() == "info": logger.setLevel(logging.INFO) + httpx_logger.setLevel(logging.INFO)