Better handle anonymization (#17207)

* Better handle anonymization

* Fix tests

Author: hxomer

litellm/proxy/guardrails/guardrail_hooks/aim/aim.py

@@ -117,7 +117,7 @@ async def call_aim_guardrail(
             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 @@ def _handle_block_action(self, analysis_result: Any, required_action: Any) -> No
         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 @@ async def call_aim_guardrail_on_output(
             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 @@ def _handle_block_action_on_output(
         )
         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 @@ async def forward_the_stream_to_aim(
             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

@@ -443,23 +443,23 @@ async def connect_mock(*args, **kwargs):
         "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": {},
             },
         },
     },