fixes for presidio checks

This commit is contained in:
Ishaan Jaff 2025-04-21 20:06:07 -07:00
parent 4a50cf10fb
commit c3e52e142e

View file

@ -238,10 +238,13 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
tasks = [] tasks = []
for m in messages: for m in messages:
if isinstance(m["content"], str): content = m.get("content", None)
if content is None:
continue
if isinstance(content, str):
tasks.append( tasks.append(
self.check_pii( self.check_pii(
text=m["content"], text=content,
output_parse_pii=self.output_parse_pii, output_parse_pii=self.output_parse_pii,
presidio_config=presidio_config, presidio_config=presidio_config,
request_data=data, request_data=data,
@ -249,7 +252,10 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
) )
responses = await asyncio.gather(*tasks) responses = await asyncio.gather(*tasks)
for index, r in enumerate(responses): for index, r in enumerate(responses):
if isinstance(messages[index]["content"], str): content = messages[index].get("content", None)
if content is None:
continue
if isinstance(content, str):
messages[index][ messages[index][
"content" "content"
] = r # replace content with redacted string ] = r # replace content with redacted string
@ -314,10 +320,11 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
for m in messages: for m in messages:
text_str = "" text_str = ""
if m["content"] is None: content = m.get("content", None)
if content is None:
continue continue
if isinstance(m["content"], str): if isinstance(content, str):
text_str = m["content"] text_str = content
tasks.append( tasks.append(
self.check_pii( self.check_pii(
text=text_str, text=text_str,
@ -328,7 +335,10 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
) # need to pass separately b/c presidio has context window limits ) # need to pass separately b/c presidio has context window limits
responses = await asyncio.gather(*tasks) responses = await asyncio.gather(*tasks)
for index, r in enumerate(responses): for index, r in enumerate(responses):
if isinstance(messages[index]["content"], str): content = messages[index].get("content", None)
if content is None:
continue
if isinstance(content, str):
messages[index][ messages[index][
"content" "content"
] = r # replace content with redacted string ] = r # replace content with redacted string
@ -373,7 +383,9 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
self, data: dict self, data: dict
) -> Optional[PresidioPerRequestConfig]: ) -> Optional[PresidioPerRequestConfig]:
if "metadata" in data: if "metadata" in data:
_metadata = data["metadata"] _metadata = data.get("metadata", None)
if _metadata is None:
return None
_guardrail_config = _metadata.get("guardrail_config") _guardrail_config = _metadata.get("guardrail_config")
if _guardrail_config: if _guardrail_config:
_presidio_config = PresidioPerRequestConfig(**_guardrail_config) _presidio_config = PresidioPerRequestConfig(**_guardrail_config)