scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-05 07:37:50 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 360

test(guardrails): key the test guardrail hook payloads by str

Browse source
This commit is contained in:
yucheng 2026-10-26 18:42:19 -07:00
parent e2304e63de
commit 0676e6a460
2 changed files with 9 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
tests/unit/proxy/policy_engine/test_pipeline_executor.py
View file

@ -1457,7 +1457,7 @@ class _LitellmMetadataOpeningGuardrail(CustomGuardrail):
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
request_data: dict[str, object],
input_type: Literal["request", "response"],
logging_obj: LiteLLMLoggingObj | None = None,
) -> GenericGuardrailAPIInputs:

16
tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py
View file

@ -1171,13 +1171,13 @@ class _PipelineStepGuardrail(CustomGuardrail):
self.calls = 0
async def async_pre_call_hook(
self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict, call_type: CallTypesLiteral
) -> dict:
self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict[str, object], call_type: CallTypesLiteral
) -> dict[str, object]:
self.calls += 1
return data
async def async_post_call_success_hook(
self, data: dict, user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes
self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes
) -> None:
self.calls += 1
@ -1252,8 +1252,8 @@ class _RaisingPipelineStepGuardrail(_PipelineStepGuardrail):
self.error = error
async def async_pre_call_hook(
self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict, call_type: CallTypesLiteral
) -> dict:
self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict[str, object], call_type: CallTypesLiteral
) -> dict[str, object]:
self.calls += 1
raise self.error
@ -1312,8 +1312,8 @@ async def test_deployment_hook_still_runs_a_pipeline_step_that_raised(
class _OpensLitellmMetadataGuardrail(_PipelineStepGuardrail):
async def async_pre_call_hook(
self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict, call_type: CallTypesLiteral
) -> dict:
self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict[str, object], call_type: CallTypesLiteral
) -> dict[str, object]:
self.calls += 1
data["litellm_metadata"] = {}
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
@ -1366,7 +1366,7 @@ async def test_post_call_pipeline_steps_are_named_only_once_they_run(
class _OpensLitellmMetadataOnResponseGuardrail(_PipelineStepGuardrail):
async def async_post_call_success_hook(
self, data: dict, user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes
self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes
) -> None:
self.calls += 1
data["litellm_metadata"] = {"user_api_key_hash": "hashed"}