fix(llm_guard.py): working llm-guard 'key-specific' mode

This commit is contained in:
Krrish Dholakia 2024-03-26 17:47:20 -07:00
parent e10eb8f6fe
commit 6d418a2920
3 changed files with 11 additions and 4 deletions

View file

@ -107,16 +107,21 @@ class _ENTERPRISE_LLMGuard(CustomLogger):
- Use the sanitized prompt returned
- LLM Guard can handle things like PII Masking, etc.
"""
self.print_verbose(f"Inside LLM Guard Pre-Call Hook")
self.print_verbose(
f"Inside LLM Guard Pre-Call Hook - llm_guard_mode={self.llm_guard_mode}"
)
# check if llm guard enabled for specific keys only
if self.llm_guard_mode == "key-specific":
self.print_verbose(
f"user_api_key_dict.permissions: {user_api_key_dict.permissions}"
)
if (
user_api_key_dict.permissions.get("enable_llm_guard_check", False)
== False
):
return
self.print_verbose("Makes LLM Guard Check")
try:
assert call_type in [
"completion",