fix(presidio_pii_masking.py): add support for setting 'http://' if unset by render env for presidio base url

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:

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