diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index c3da689220..c0903a35b6 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -12,6 +12,7 @@ from litellm._uuid import uuid from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy.utils import PrismaClient +from litellm.proxy.types_utils.utils import get_instance_fn from litellm.secret_managers.main import get_secret from litellm.types.guardrails import ( Guardrail, @@ -489,7 +490,7 @@ class InMemoryGuardrailHandler: config_file_path: Optional[str] = None, ) -> Optional[CustomGuardrail]: """ - Initialize a Custom Guardrail from a python file + Initialize a Custom Guardrail from a python file or module path This initializes it by adding it to the litellm callback manager """ @@ -498,26 +499,12 @@ class InMemoryGuardrailHandler: "GuardrailsAIException - Please pass the config_file_path to initialize_guardrails_v2" ) - _file_name, _class_name = guardrail_type.split(".") verbose_proxy_logger.debug( - "Initializing custom guardrail: %s, file_name: %s, class_name: %s", + "Initializing custom guardrail: %s", guardrail_type, - _file_name, - _class_name, ) - directory = os.path.dirname(config_file_path) - module_file_path = os.path.join(directory, _file_name) + ".py" - - spec = importlib.util.spec_from_file_location(_class_name, module_file_path) # type: ignore - if not spec: - raise ImportError( - f"Could not find a module specification for {module_file_path}" - ) - - module = importlib.util.module_from_spec(spec) # type: ignore - spec.loader.exec_module(module) # type: ignore - _guardrail_class = getattr(module, _class_name) + _guardrail_class = get_instance_fn(guardrail_type, config_file_path=config_file_path) mode = litellm_params.mode if mode is None: