phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

fix(presidio_pii_masking.py): fix presidio unset url check + add same check for langfuse

Browse source
This commit is contained in:
Krrish Dholakia 2024-07-06 17:50:55 -07:00
parent d57d3df1d6
commit 1193ee8803
4 changed files with 84 additions and 69 deletions
Download patch file Download diff file Expand all files Collapse all files

6
litellm/integrations/langfuse.py
View file

@ -32,6 +32,12 @@ class LangFuseLogger:
self.langfuse_host = langfuse_host or os.getenv(
"LANGFUSE_HOST", "https://cloud.langfuse.com"
)
if not (
self.langfuse_host.startswith("http://")
or self.langfuse_host.startswith("https://")
):
# add http:// if unset, assume communicating over private network - e.g. render
self.langfuse_host = "http://" + self.langfuse_host
self.langfuse_release = os.getenv("LANGFUSE_RELEASE")
self.langfuse_debug = os.getenv("LANGFUSE_DEBUG")

16
litellm/proxy/hooks/presidio_pii_masking.py
View file

@ -70,15 +70,16 @@ class _OPTIONAL_PresidioPIIMasking(CustomLogger):
) # type: ignore
self.presidio_anonymizer_api_base: Optional[str] = litellm.get_secret(
"PRESIDIO_ANONYMIZER_API_BASE", None
)
) # type: ignore
if self.presidio_analyzer_api_base is None:
raise Exception("Missing `PRESIDIO_ANALYZER_API_BASE` from environment")
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://"):
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
@ -88,9 +89,10 @@ class _OPTIONAL_PresidioPIIMasking(CustomLogger):
raise Exception("Missing `PRESIDIO_ANONYMIZER_API_BASE` from environment")
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://"):
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

101
litellm/tests/langfuse.log
View file

File diff suppressed because one or more lines are too long

30
litellm/tests/test_presidio_masking.py
View file

@ -26,23 +26,31 @@ from litellm.proxy.hooks.presidio_pii_masking import _OPTIONAL_PresidioPIIMaskin
from litellm.proxy.utils import ProxyLogging
def test_validate_environment_missing_http():
@pytest.mark.parametrize(
"base_url",
[
"presidio-analyzer-s3pa:10000",
"https://presidio-analyzer-s3pa:10000",
"http://presidio-analyzer-s3pa:10000",
],
)
def test_validate_environment_missing_http(base_url):
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"
)
os.environ["PRESIDIO_ANALYZER_API_BASE"] = f"{base_url}/analyze"
os.environ["PRESIDIO_ANONYMIZER_API_BASE"] = f"{base_url}/anonymize"
pii_masking.validate_environment()
expected_url = base_url
if not (base_url.startswith("https://") or base_url.startswith("http://")):
expected_url = "http://" + base_url
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/"
pii_masking.presidio_anonymizer_api_base == f"{expected_url}/anonymize/"
), "Got={}, Expected={}".format(
pii_masking.presidio_anonymizer_api_base, f"{expected_url}/anonymize/"
)
assert pii_masking.presidio_analyzer_api_base == f"{expected_url}/analyze/"
@pytest.mark.asyncio