From d30de53fb8596c09e22fef361f89c60f4b66866b Mon Sep 17 00:00:00 2001 From: Josephasafg Date: Thu, 20 Jun 2024 16:02:11 +0300 Subject: [PATCH 1/9] fix: Added logger to httpx --- ai21/logger.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/ai21/logger.py b/ai21/logger.py index 49621c03..beb3fa23 100644 --- a/ai21/logger.py +++ b/ai21/logger.py @@ -3,6 +3,7 @@ from ai21.ai21_env_config import AI21EnvConfig logger = logging.getLogger("ai21") +httpx_logger = logging.getLogger("httpx") def _basic_config() -> None: @@ -17,5 +18,7 @@ def setup_logger() -> None: if AI21EnvConfig.log_level == "debug": logger.setLevel(logging.DEBUG) + httpx_logger.setLevel(logging.DEBUG) elif AI21EnvConfig.log_level == "info": logger.setLevel(logging.INFO) + httpx_logger.setLevel(logging.INFO) From d79462d0ed1ca54e4409a24dc5394cd0c8ce16af Mon Sep 17 00:00:00 2001 From: Josephasafg Date: Mon, 24 Jun 2024 10:37:58 +0300 Subject: [PATCH 2/9] feat: Added set_verbose --- ai21/http_client/async_http_client.py | 2 +- ai21/http_client/base_http_client.py | 17 +++++++++----- ai21/http_client/http_client.py | 2 +- ai21/logger.py | 32 +++++++++++++++++++++++++-- 4 files changed, 44 insertions(+), 9 deletions(-) diff --git a/ai21/http_client/async_http_client.py b/ai21/http_client/async_http_client.py index 914df7cd..6943c99f 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}") + self._log_request(method=method, url=url, headers=headers, params=params, body=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..5e8cd693 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,13 @@ ServiceUnavailable, AI21APIError, ) +from ai21.logger import logger, get_verbose +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 @@ -116,3 +115,11 @@ def _request( @abstractmethod def _init_client(self, client: Optional[_HttpxClientT]) -> _HttpxClientT: pass + + def _log_request( + self, method: str, url: str, headers: Dict[str, Any], params: Dict[str, Any], body: Dict[str, Any] + ) -> None: + if not get_verbose(): + headers = {key: value for key, value in headers.items() if key != "api_key" and key != "Authorization"} + + logger.debug(f"Calling {method} {url} {headers} {params} {body}") diff --git a/ai21/http_client/http_client.py b/ai21/http_client/http_client.py index b7500806..fca16a17 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}") + self._log_request(method=method, url=url, headers=headers, params=params, body=body) if method == "GET": request = self._client.build_request( diff --git a/ai21/logger.py b/ai21/logger.py index beb3fa23..0175c257 100644 --- a/ai21/logger.py +++ b/ai21/logger.py @@ -1,11 +1,34 @@ +import os import logging from ai21.ai21_env_config import AI21EnvConfig +_verbose = False + logger = logging.getLogger("ai21") httpx_logger = logging.getLogger("httpx") +def set_verbose(value: bool) -> None: + """ " + Use this function if you want to log additional, more sensitive data like - secrets and environment variables + """ + global _verbose + _verbose = value + + if _verbose: + 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: logging.basicConfig( format="[%(asctime)s - %(name)s - %(levelname)s] %(message)s", @@ -13,12 +36,17 @@ def _basic_config() -> None: ) +def log_verbose(text: str) -> None: + if get_verbose(): + logger.info(text) + + def setup_logger() -> None: _basic_config() - if AI21EnvConfig.log_level == "debug": + if AI21EnvConfig.log_level.lower() == "debug": logger.setLevel(logging.DEBUG) httpx_logger.setLevel(logging.DEBUG) - elif AI21EnvConfig.log_level == "info": + elif AI21EnvConfig.log_level.lower() == "info": logger.setLevel(logging.INFO) httpx_logger.setLevel(logging.INFO) From 6812413f7198972b83dcf6702c9615fb77316c99 Mon Sep 17 00:00:00 2001 From: Josephasafg Date: Mon, 24 Jun 2024 11:20:35 +0300 Subject: [PATCH 3/9] feat: Added set_debug --- ai21/__init__.py | 3 ++- ai21/logger.py | 16 ++++++++++++---- 2 files changed, 14 insertions(+), 5 deletions(-) diff --git a/ai21/__init__.py b/ai21/__init__.py index bf6eda84..97ff412b 100644 --- a/ai21/__init__.py +++ b/ai21/__init__.py @@ -18,7 +18,6 @@ from ai21.version import VERSION __version__ = VERSION -setup_logger() def _import_bedrock_client(): @@ -70,3 +69,5 @@ def __getattr__(name: str) -> Any: "AI21AzureClient", "AsyncAI21AzureClient", ] + +setup_logger() diff --git a/ai21/logger.py b/ai21/logger.py index 0175c257..2780b9c9 100644 --- a/ai21/logger.py +++ b/ai21/logger.py @@ -11,15 +11,23 @@ def set_verbose(value: bool) -> None: """ " - Use this function if you want to log additional, more sensitive data like - secrets and environment variables + 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 - if _verbose: - os.environ["AI21_LOG_LEVEL"] = "DEBUG" + set_debug(_verbose) + + +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" + os.environ["AI21_LOG_LEVEL"] = "info" setup_logger() From 6a64ab0ca7479c9900aaa1592dee8db08183aa29 Mon Sep 17 00:00:00 2001 From: Josephasafg Date: Mon, 24 Jun 2024 11:29:43 +0300 Subject: [PATCH 4/9] fix: api-key header --- ai21/http_client/base_http_client.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ai21/http_client/base_http_client.py b/ai21/http_client/base_http_client.py index 5e8cd693..945ff978 100644 --- a/ai21/http_client/base_http_client.py +++ b/ai21/http_client/base_http_client.py @@ -120,6 +120,6 @@ def _log_request( self, method: str, url: str, headers: Dict[str, Any], params: Dict[str, Any], body: Dict[str, Any] ) -> None: if not get_verbose(): - headers = {key: value for key, value in headers.items() if key != "api_key" and key != "Authorization"} + headers = {key: value for key, value in headers.items() if key != "api-key" and key != "Authorization"} logger.debug(f"Calling {method} {url} {headers} {params} {body}") From 99a6fd444351cb2ecf6533f48795f7a464da6ec7 Mon Sep 17 00:00:00 2001 From: Josephasafg Date: Mon, 24 Jun 2024 13:06:48 +0300 Subject: [PATCH 5/9] fix: Removed unused function --- ai21/logger.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/ai21/logger.py b/ai21/logger.py index 2780b9c9..5bfd052a 100644 --- a/ai21/logger.py +++ b/ai21/logger.py @@ -44,11 +44,6 @@ def _basic_config() -> None: ) -def log_verbose(text: str) -> None: - if get_verbose(): - logger.info(text) - - def setup_logger() -> None: _basic_config() From e647583bda88b64093bce87b390d3e85f6f27b43 Mon Sep 17 00:00:00 2001 From: Josephasafg Date: Mon, 24 Jun 2024 13:18:14 +0300 Subject: [PATCH 6/9] feat: Logged env variables --- ai21/ai21_env_config.py | 19 +++++++++++++++++++ ai21/http_client/base_http_client.py | 6 +++++- 2 files changed, 24 insertions(+), 1 deletion(-) 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/base_http_client.py b/ai21/http_client/base_http_client.py index 945ff978..acbf4d57 100644 --- a/ai21/http_client/base_http_client.py +++ b/ai21/http_client/base_http_client.py @@ -6,6 +6,7 @@ import httpx +from ai21 import AI21EnvConfig from ai21.errors import ( BadRequest, Unauthorized, @@ -119,7 +120,10 @@ def _init_client(self, client: Optional[_HttpxClientT]) -> _HttpxClientT: def _log_request( self, method: str, url: str, headers: Dict[str, Any], params: Dict[str, Any], body: Dict[str, Any] ) -> None: - if not get_verbose(): + is_verbose = get_verbose() + + if not is_verbose: headers = {key: value for key, value in headers.items() if key != "api-key" and key != "Authorization"} logger.debug(f"Calling {method} {url} {headers} {params} {body}") + AI21EnvConfig.log(with_secrets=is_verbose) From 502f9c0b279084bef5977870e7d7deb3887e5087 Mon Sep 17 00:00:00 2001 From: Josephasafg Date: Mon, 24 Jun 2024 13:33:30 +0300 Subject: [PATCH 7/9] fix: Changed call location --- ai21/__init__.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/ai21/__init__.py b/ai21/__init__.py index 97ff412b..bf6eda84 100644 --- a/ai21/__init__.py +++ b/ai21/__init__.py @@ -18,6 +18,7 @@ from ai21.version import VERSION __version__ = VERSION +setup_logger() def _import_bedrock_client(): @@ -69,5 +70,3 @@ def __getattr__(name: str) -> Any: "AI21AzureClient", "AsyncAI21AzureClient", ] - -setup_logger() From 89a104b330e8dd1c4fe0701c3d4f6a830d21d321 Mon Sep 17 00:00:00 2001 From: Josephasafg Date: Sun, 30 Jun 2024 11:40:03 +0300 Subject: [PATCH 8/9] fix: CR --- ai21/http_client/async_http_client.py | 2 +- ai21/http_client/base_http_client.py | 13 ------------ ai21/http_client/http_client.py | 2 +- ai21/logger.py | 30 +++++++++++++++++++++++++-- 4 files changed, 30 insertions(+), 17 deletions(-) diff --git a/ai21/http_client/async_http_client.py b/ai21/http_client/async_http_client.py index 6943c99f..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 - self._log_request(method=method, url=url, headers=headers, params=params, body=body) + 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 acbf4d57..bf9efe98 100644 --- a/ai21/http_client/base_http_client.py +++ b/ai21/http_client/base_http_client.py @@ -6,7 +6,6 @@ import httpx -from ai21 import AI21EnvConfig from ai21.errors import ( BadRequest, Unauthorized, @@ -16,7 +15,6 @@ ServiceUnavailable, AI21APIError, ) -from ai21.logger import logger, get_verbose from ai21.stream.async_stream import AsyncStream from ai21.stream.stream import Stream @@ -116,14 +114,3 @@ def _request( @abstractmethod def _init_client(self, client: Optional[_HttpxClientT]) -> _HttpxClientT: pass - - def _log_request( - self, method: str, url: str, headers: Dict[str, Any], params: Dict[str, Any], body: Dict[str, Any] - ) -> None: - is_verbose = get_verbose() - - if not is_verbose: - headers = {key: value for key, value in headers.items() if key != "api-key" and key != "Authorization"} - - logger.debug(f"Calling {method} {url} {headers} {params} {body}") - AI21EnvConfig.log(with_secrets=is_verbose) diff --git a/ai21/http_client/http_client.py b/ai21/http_client/http_client.py index fca16a17..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 - self._log_request(method=method, url=url, headers=headers, params=params, body=body) + 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 5bfd052a..d51f46b1 100644 --- a/ai21/logger.py +++ b/ai21/logger.py @@ -1,5 +1,6 @@ -import os import logging +import os +import re from ai21.ai21_env_config import AI21EnvConfig @@ -9,8 +10,29 @@ 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*'[^']*')" + + 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. """ @@ -19,6 +41,8 @@ def set_verbose(value: bool) -> None: set_debug(_verbose) + AI21EnvConfig.log(with_secrets=value) + def set_debug(value: bool) -> None: """ @@ -46,6 +70,8 @@ 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.lower() == "debug": logger.setLevel(logging.DEBUG) From 81c9d93b0e5b2aa1090064420df0828249ae75e0 Mon Sep 17 00:00:00 2001 From: Josephasafg Date: Mon, 1 Jul 2024 13:43:19 +0300 Subject: [PATCH 9/9] fix: Added amazon header to secrets --- ai21/logger.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ai21/logger.py b/ai21/logger.py index d51f46b1..a0706e5b 100644 --- a/ai21/logger.py +++ b/ai21/logger.py @@ -22,7 +22,7 @@ def format(self, record: logging.LogRecord) -> str: def _censor_secrets(self, message: str) -> str: # Regular expression to find the Authorization key and its value - pattern = r"('Authorization':\s*'[^']*'|'api-key':\s*'[^']*')" + pattern = r"('Authorization':\s*'[^']*'|'api-key':\s*'[^']*'|'X-Amz-Security-Token':\s*'[^']*')" def replacement(match): return match.group(0).split(":")[0] + ": '**************'"