diff --git a/connectchain/lcel/logger.py b/connectchain/lcel/logger.py index fdff934..8b85725 100644 --- a/connectchain/lcel/logger.py +++ b/connectchain/lcel/logger.py @@ -20,10 +20,15 @@ class Logger(ABC): def print(self, payload): """Send the payload to the desired component""" + @abstractmethod + def sanitize_output(self, payload): + """Sanitize the payload to remove sensitive information""" + def log(self): """Return a lambda function that logs the payload""" def _log_helper(payload: any): + self.sanitize_output(payload) self.print(payload) return payload diff --git a/connectchain/test/test_lcel_logger.py b/connectchain/test/test_lcel_logger.py index 8ad5785..4aea1a4 100644 --- a/connectchain/test/test_lcel_logger.py +++ b/connectchain/test/test_lcel_logger.py @@ -10,6 +10,7 @@ # or implied. See the License for the specific language governing permissions and limitations under # the License. """This is the unit test for the lcel_logger""" +import re from unittest import TestCase from connectchain.prompts import ValidPromptTemplate @@ -23,6 +24,9 @@ def __init__(self): def print(self, payload): self.output = payload.text + def sanitize_output(self, str): + string = re.sub(r"^(?:(?:\+?1\s*(?:[.-]\s*)?)?(?:\(\s*([2-9]1[02-9]|[2-9][02-8]1|[2-9][02-8][02-9])\s*\)|([2-9]1[02-9]|[2-9][02-8]1|[2-9][02-8][02-9]))\s*(?:[.-]\s*)?)?([2-9]1[02-9]|[2-9][02-9]1|[2-9][02-9]{2})\s*(?:[.-]\s*)?([0-9]{4})(?:\s*(?:#|x\.?|ext\.?|extension)\s*(\d+))?$", "", payload) + self.output = string TEMPLATE = "Give me 10 different {animal_type} species."