diff --git a/litellm/proxy/hooks/presidio_pii_masking.py b/litellm/proxy/hooks/presidio_pii_masking.py index e64e69c4579..609fbc7d47d 100644 --- a/litellm/proxy/hooks/presidio_pii_masking.py +++ b/litellm/proxy/hooks/presidio_pii_masking.py @@ -8,21 +8,26 @@ # Tell us how we can improve! - Krrish & Ishaan +import asyncio +import json +import traceback +import uuid from typing import Optional, Union -import litellm, traceback, uuid, json # noqa: E401 -from litellm.caching import DualCache -from litellm.proxy._types import UserAPIKeyAuth -from litellm.integrations.custom_logger import CustomLogger + +import aiohttp from fastapi import HTTPException + +import litellm # noqa: E401 from litellm._logging import verbose_proxy_logger +from litellm.caching import DualCache +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy._types import UserAPIKeyAuth from litellm.utils import ( - ModelResponse, EmbeddingResponse, ImageResponse, + ModelResponse, StreamingChoices, ) -import aiohttp -import asyncio class _OPTIONAL_PresidioPIIMasking(CustomLogger): @@ -57,22 +62,39 @@ class _OPTIONAL_PresidioPIIMasking(CustomLogger): f"An error occurred: {str(e)}, file_path={ad_hoc_recognizers}" ) - self.presidio_analyzer_api_base = litellm.get_secret( + self.validate_environment() + + def validate_environment(self): + self.presidio_analyzer_api_base: Optional[str] = litellm.get_secret( "PRESIDIO_ANALYZER_API_BASE", None - ) - self.presidio_anonymizer_api_base = litellm.get_secret( + ) # type: ignore + self.presidio_anonymizer_api_base: Optional[str] = litellm.get_secret( "PRESIDIO_ANONYMIZER_API_BASE", None ) if self.presidio_analyzer_api_base is None: raise Exception("Missing `PRESIDIO_ANALYZER_API_BASE` from environment") - elif not self.presidio_analyzer_api_base.endswith("/"): + if not self.presidio_analyzer_api_base.endswith("/"): self.presidio_analyzer_api_base += "/" + if not self.presidio_analyzer_api_base.startswith( + "http://" + ) or self.presidio_analyzer_api_base.startswith("https://"): + # add http:// if unset, assume communicating over private network - e.g. render + self.presidio_analyzer_api_base = ( + "http://" + self.presidio_analyzer_api_base + ) if self.presidio_anonymizer_api_base is None: raise Exception("Missing `PRESIDIO_ANONYMIZER_API_BASE` from environment") - elif not self.presidio_anonymizer_api_base.endswith("/"): + if not self.presidio_anonymizer_api_base.endswith("/"): self.presidio_anonymizer_api_base += "/" + if not self.presidio_anonymizer_api_base.startswith( + "http://" + ) or self.presidio_anonymizer_api_base.startswith("https://"): + # add http:// if unset, assume communicating over private network - e.g. render + self.presidio_anonymizer_api_base = ( + "http://" + self.presidio_anonymizer_api_base + ) def print_verbose(self, print_statement): try: diff --git a/litellm/tests/test_presidio_masking.py b/litellm/tests/test_presidio_masking.py index f8ac9820fd6..2b02ea298d4 100644 --- a/litellm/tests/test_presidio_masking.py +++ b/litellm/tests/test_presidio_masking.py @@ -1,8 +1,13 @@ # What is this? ## Unit test for presidio pii masking -import sys, os, asyncio, time, random -from datetime import datetime +import asyncio +import os +import random +import sys +import time import traceback +from datetime import datetime + from dotenv import load_dotenv load_dotenv() @@ -12,12 +17,32 @@ sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path import pytest + import litellm -from litellm.proxy.hooks.presidio_pii_masking import _OPTIONAL_PresidioPIIMasking from litellm import Router, mock_completion -from litellm.proxy.utils import ProxyLogging -from litellm.proxy._types import UserAPIKeyAuth from litellm.caching import DualCache +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.hooks.presidio_pii_masking import _OPTIONAL_PresidioPIIMasking +from litellm.proxy.utils import ProxyLogging + + +def test_validate_environment_missing_http(): + pii_masking = _OPTIONAL_PresidioPIIMasking(mock_testing=True) + + os.environ["PRESIDIO_ANALYZER_API_BASE"] = "presidio-analyzer-s3pa:10000/analyze" + os.environ["PRESIDIO_ANONYMIZER_API_BASE"] = ( + "presidio-analyzer-s3pa:10000/anonymize" + ) + pii_masking.validate_environment() + + assert ( + pii_masking.presidio_anonymizer_api_base + == "http://presidio-analyzer-s3pa:10000/anonymize/" + ) + assert ( + pii_masking.presidio_analyzer_api_base + == "http://presidio-analyzer-s3pa:10000/analyze/" + ) @pytest.mark.asyncio