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 - Use the sanitized prompt returned
- LLM Guard can handle things like PII Masking, etc. - 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 # check if llm guard enabled for specific keys only
if self.llm_guard_mode == "key-specific": if self.llm_guard_mode == "key-specific":
self.print_verbose(
f"user_api_key_dict.permissions: {user_api_key_dict.permissions}"
)
if ( if (
user_api_key_dict.permissions.get("enable_llm_guard_check", False) user_api_key_dict.permissions.get("enable_llm_guard_check", False)
== False == False
): ):
return return
self.print_verbose("Makes LLM Guard Check")
try: try:
assert call_type in [ assert call_type in [
"completion", "completion",

View file

@ -1,6 +1,6 @@
### INIT VARIABLES ### ### INIT VARIABLES ###
import threading, requests, os import threading, requests, os
from typing import Callable, List, Optional, Dict, Union, Any from typing import Callable, List, Optional, Dict, Union, Any, Literal
from litellm.caching import Cache from litellm.caching import Cache
from litellm._logging import set_verbose, _turn_on_debug, verbose_logger from litellm._logging import set_verbose, _turn_on_debug, verbose_logger
from litellm.proxy._types import KeyManagementSystem, KeyManagementSettings from litellm.proxy._types import KeyManagementSystem, KeyManagementSettings

View file

@ -6,7 +6,7 @@ from datetime import datetime
import importlib import importlib
from dotenv import load_dotenv from dotenv import load_dotenv
import urllib.parse as urlparse import urllib.parse as urlparse
from litellm._logging import verbose_proxy_logger
sys.path.append(os.getcwd()) sys.path.append(os.getcwd())
@ -20,6 +20,8 @@ telemetry = None
def append_query_params(url, params): def append_query_params(url, params):
from litellm._logging import verbose_proxy_logger
verbose_proxy_logger.debug(f"url: {url}") verbose_proxy_logger.debug(f"url: {url}")
verbose_proxy_logger.debug(f"params: {params}") verbose_proxy_logger.debug(f"params: {params}")
parsed_url = urlparse.urlparse(url) parsed_url = urlparse.urlparse(url)