mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
fix(xecguard): use StandardLoggingGuardrailInformation in logging hook (#32911)
XecGuard's async_logging_hook wrote a bare dict to standard_logging_object["guardrail_information"] while the typed contract is Optional[List[StandardLoggingGuardrailInformation]]. Readers that iterated the field walked dict keys, raised on info.get, or silently dropped the entry from guardrail usage tracking and spend-log writes Construct the typed entry and append it to the existing list or create a new one, matching the shared helper pattern. Record the configured guardrail name instead of a hardcoded "xecguard" and pass the GuardrailEventHooks enum for guardrail_mode
This commit is contained in:
parent
81db114c40
commit
f947ef14a2
2 changed files with 50 additions and 14 deletions
|
|
@ -49,7 +49,11 @@ from litellm.llms.custom_httpx.http_handler import (
|
|||
httpxSpecialProvider,
|
||||
)
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs, GuardrailStatus
|
||||
from litellm.types.utils import (
|
||||
GenericGuardrailAPIInputs,
|
||||
GuardrailStatus,
|
||||
StandardLoggingGuardrailInformation,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
|
|
@ -246,16 +250,21 @@ class XecGuardGuardrail(CustomGuardrail):
|
|||
"guardrail_intervened" if scan_result.get("decision") == "UNSAFE" else "success"
|
||||
)
|
||||
end_time = datetime.now()
|
||||
kwargs["standard_logging_object"]["guardrail_information"] = {
|
||||
"duration": (end_time - start_time).total_seconds(),
|
||||
"end_time": end_time.timestamp(),
|
||||
"guardrail_mode": "logging_only",
|
||||
"guardrail_name": "xecguard",
|
||||
"guardrail_response": scan_result,
|
||||
"guardrail_status": guardrail_status,
|
||||
"masked_entity_count": None,
|
||||
"start_time": start_time.timestamp(),
|
||||
}
|
||||
slg = StandardLoggingGuardrailInformation(
|
||||
guardrail_name=self.guardrail_name or "xecguard",
|
||||
guardrail_mode=GuardrailEventHooks.logging_only,
|
||||
guardrail_response=scan_result,
|
||||
guardrail_status=guardrail_status,
|
||||
start_time=start_time.timestamp(),
|
||||
end_time=end_time.timestamp(),
|
||||
duration=(end_time - start_time).total_seconds(),
|
||||
masked_entity_count=None,
|
||||
)
|
||||
existing = kwargs["standard_logging_object"].get("guardrail_information")
|
||||
if isinstance(existing, list):
|
||||
existing.append(slg)
|
||||
else:
|
||||
kwargs["standard_logging_object"]["guardrail_information"] = [slg]
|
||||
|
||||
except Exception as exc:
|
||||
verbose_proxy_logger.debug(
|
||||
|
|
|
|||
|
|
@ -1644,12 +1644,37 @@ class TestXecGuardLoggingHook:
|
|||
)
|
||||
assert out_kwargs is kwargs
|
||||
assert out_result is result
|
||||
info = kwargs["standard_logging_object"]["guardrail_information"]
|
||||
info_list = kwargs["standard_logging_object"]["guardrail_information"]
|
||||
assert isinstance(info_list, list), "guardrail_information must be a list"
|
||||
assert len(info_list) == 1
|
||||
info = info_list[0]
|
||||
assert info["guardrail_mode"] == "logging_only"
|
||||
assert info["guardrail_name"] == "xecguard"
|
||||
assert info["guardrail_name"] == "test-xecguard"
|
||||
assert info["guardrail_status"] == "success"
|
||||
assert info["guardrail_response"]["trace_id"] == "lg-1"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_logging_hook_appends_to_existing_guardrail_info(
|
||||
self, xecguard_guardrail, mock_request_data
|
||||
):
|
||||
resp = _make_response({"decision": "SAFE", "trace_id": "lg-4"})
|
||||
prior_entry = {"guardrail_name": "other-guardrail"}
|
||||
with patch.object(xecguard_guardrail.async_handler, "post", return_value=resp):
|
||||
kwargs = {
|
||||
**mock_request_data,
|
||||
"standard_logging_object": {"guardrail_information": [prior_entry]},
|
||||
}
|
||||
await xecguard_guardrail.async_logging_hook(
|
||||
kwargs=kwargs,
|
||||
result=_build_model_response("some answer"),
|
||||
call_type="acompletion",
|
||||
)
|
||||
info_list = kwargs["standard_logging_object"]["guardrail_information"]
|
||||
assert len(info_list) == 2
|
||||
assert info_list[0] is prior_entry
|
||||
assert info_list[1]["guardrail_name"] == "test-xecguard"
|
||||
assert info_list[1]["guardrail_response"]["trace_id"] == "lg-4"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_logging_hook_without_response_records_info(
|
||||
self, xecguard_guardrail, mock_request_data
|
||||
|
|
@ -1680,7 +1705,9 @@ class TestXecGuardLoggingHook:
|
|||
result=_build_model_response("x"),
|
||||
call_type="acompletion",
|
||||
)
|
||||
info = kwargs["standard_logging_object"]["guardrail_information"]
|
||||
info_list = kwargs["standard_logging_object"]["guardrail_information"]
|
||||
assert isinstance(info_list, list), "guardrail_information must be a list"
|
||||
info = info_list[0]
|
||||
assert info["guardrail_status"] == "guardrail_intervened"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue