fix(utils.py): handle key error in msg validation (#8325)

* fix(utils.py): handle key error in msg validation

* Support running Aim Guard during LLM call (#7918)

* support running Aim Guard during LLM call

* Rename header

* adjust docs and fix type annotations

* fix(timeout.md): doc fix for openai example on dynamic timeouts

---------

Co-authored-by: Tomer Bin <117278227+hxtomer@users.noreply.github.com>
This commit is contained in:
Krish Dholakia 2025-02-06 18:13:46 -08:00 committed by GitHub
parent fac1d2ccef
commit f031926b82
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 56 additions and 21 deletions

View file

@ -55,15 +55,15 @@ def test_aim_guard_config_no_api_key():
@pytest.mark.asyncio
async def test_callback():
@pytest.mark.parametrize("mode", ["pre_call", "during_call"])
async def test_callback(mode: str):
init_guardrails_v2(
all_guardrails=[
{
"guardrail_name": "gibberish-guard",
"litellm_params": {
"guardrail": "aim",
"guard_name": "gibberish_guard",
"mode": "pre_call",
"mode": mode,
"api_key": "hs-aim-key",
},
}
@ -89,6 +89,11 @@ async def test_callback():
request=Request(method="POST", url="http://aim"),
),
):
await aim_guardrail.async_pre_call_hook(
data=data, cache=DualCache(), user_api_key_dict=UserAPIKeyAuth(), call_type="completion"
)
if mode == "pre_call":
await aim_guardrail.async_pre_call_hook(
data=data, cache=DualCache(), user_api_key_dict=UserAPIKeyAuth(), call_type="completion"
)
else:
await aim_guardrail.async_moderation_hook(
data=data, user_api_key_dict=UserAPIKeyAuth(), call_type="completion"
)