tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-11 00:25:15 +00:00
Code Issues Packages Projects Releases Wiki Activity

Better handle anonymization (#17207)

Browse Source 
* Better handle anonymization

* Fix tests
This commit is contained in:
hxomer
2025-11-28 08:22:19 -08:00
committed by GitHub
parent bcc35a6069
commit b85df0b1fb
2 changed files with 22 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/guardrails/guardrail_hooks/aim/aim.py
+10 -27
View File
@@ -117,7 +117,7 @@ class AimGuardrail(CustomGuardrail):
self._handle_block_action(res["analysis_result"], required_action)
elif action_type == "anonymize_action":
return self._anonymize_request(
res["analysis_result"], required_action, data
res, data
)
else:
verbose_proxy_logger.error(f"Aim: {action_type} action")
@@ -133,27 +133,18 @@ class AimGuardrail(CustomGuardrail):
raise HTTPException(status_code=400, detail=detection_message)
def _anonymize_request(
self, analysis_result: Any, required_action: Any, data: dict
self, res: Any, data: dict
) -> dict:
verbose_proxy_logger.info("Aim: anonymize action")
redaction_result = required_action and required_action.get(
"chat_redaction_result"
)
if not redaction_result:
redacted_chat = res.get("redacted_chat")
if not redacted_chat:
return data
if analysis_result and analysis_result.get("session_entities"):
self._set_dlp_entities(analysis_result.get("session_entities"))
data["messages"] = [
{
"role": redaction_result["redacted_new_message"]["role"],
"content": redaction_result["redacted_new_message"]["content"],
}
] + [
{
"role": message["role"],
"content": message["content"],
}
for message in redaction_result["all_redacted_messages"]
for message in redacted_chat["all_redacted_messages"]
]
return data
@@ -185,7 +176,11 @@ class AimGuardrail(CustomGuardrail):
return self._handle_block_action_on_output(
res["analysis_result"], required_action
)
return self._deanonymize_output(output)
redacted_chat = res.get("redacted_chat", None)
if action_type and action_type == "anonymize_action" and redacted_chat:
return {"redacted_output": redacted_chat["all_redacted_messages"][-1]["content"]}
return {"redacted_output": output}
def _handle_block_action_on_output(
self, analysis_result: Any, required_action: Any
@@ -199,15 +194,6 @@ class AimGuardrail(CustomGuardrail):
)
return {"detection_message": detection_message}
def _deanonymize_output(self, output: str) -> dict | None:
try:
for entity in self.dlp_entities:
output = output.replace(f"[{entity['name']}]", entity["content"])
return {"redacted_output": output}
except Exception as e:
verbose_proxy_logger.error(f"Aim: Error while redacting output: {e}")
return None
def _build_aim_headers(
self,
*,
@@ -323,9 +309,6 @@ class AimGuardrail(CustomGuardrail):
await websocket.send(chunk)
await websocket.send(json.dumps({"done": True}))
def _set_dlp_entities(self, entities: list[dict]) -> None:
self.dlp_entities = entities[: self._max_dlp_entities]
@staticmethod
def get_config_model() -> Optional[Type["GuardrailConfigModel"]]:
from litellm.types.proxy.guardrails.guardrail_hooks.aim import (
tests/local_testing/test_aim_guardrails.py
+12 -12
View File
@@ -443,23 +443,23 @@ response_with_detections = Response(
"required_action": {
"action_type": "anonymize_action",
"policy_name": "PII",
"chat_redaction_result": {
"all_redacted_messages": [
{
"content": "Hi my name is [NAME_1]",
"role": "user",
"additional_contents": [],
"received_message_id": "0",
"extra_fields": {},
}
],
"redacted_new_message": {
},
"redacted_chat": {
"all_redacted_messages": [
{
"content": "Hi my name is [NAME_1]",
"role": "user",
"additional_contents": [],
"received_message_id": "0",
"extra_fields": {},
},
}
],
"redacted_new_message": {
"content": "Hi my name is [NAME_1]",
"role": "user",
"additional_contents": [],
"received_message_id": "0",
"extra_fields": {},
},
},
},