From 1b1e0f2d77699a0bd6aef3e6d05a77ed242d2845 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Fri, 23 Aug 2024 10:54:42 -0700
Subject: [PATCH] init custom guardrail class

---
 litellm/proxy/guardrails/init_guardrails.py | 25 +++++++++++++++++++--
 litellm/proxy/proxy_server.py               |  4 +++-
 2 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/litellm/proxy/guardrails/init_guardrails.py b/litellm/proxy/guardrails/init_guardrails.py
index 4180bce2d38..643e135961d 100644
--- a/litellm/proxy/guardrails/init_guardrails.py
+++ b/litellm/proxy/guardrails/init_guardrails.py
@@ -84,7 +84,10 @@ Map guardrail_name: <pre_call>, <post_call>, during_call
 """
 
 
-def init_guardrails_v2(all_guardrails: dict):
+def init_guardrails_v2(
+    all_guardrails: dict,
+    config_file_path: str,
+):
     # Convert the loaded data to the TypedDict structure
     guardrail_list = []
 
@@ -166,6 +169,10 @@ def init_guardrails_v2(all_guardrails: dict):
             isinstance(litellm_params["guardrail"], str)
             and "." in litellm_params["guardrail"]
         ):
+            import os
+
+            from litellm.proxy.utils import get_instance_fn
+
             # Custom guardrail
             _guardrail = litellm_params["guardrail"]
             _file_name, _class_name = _guardrail.split(".")
@@ -175,7 +182,21 @@ def init_guardrails_v2(all_guardrails: dict):
                 _file_name,
                 _class_name,
             )
-            _guardrail_class = getattr(importlib.import_module(_file_name), _class_name)
+
+            directory = os.path.dirname(config_file_path)
+            module_file_path = os.path.join(directory, _file_name)
+            module_file_path += ".py"
+
+            spec = importlib.util.spec_from_file_location(_class_name, module_file_path)  # type: ignore
+            if spec is None:
+                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_callback = _guardrail_class(
                 guardrail_name=guardrail["guardrail_name"],
                 event_hook=litellm_params["mode"],
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 3ef5609db37..f4206f726ac 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -1959,7 +1959,9 @@ class ProxyConfig:
         # Guardrail settings
         guardrails_v2 = config.get("guardrails", None)
         if guardrails_v2:
-            init_guardrails_v2(all_guardrails=guardrails_v2)
+            init_guardrails_v2(
+                all_guardrails=guardrails_v2, config_file_path=config_file_path
+            )
         return router, router.get_model_list(), general_settings
 
     def get_model_info_with_id(self, model, db_model=False) -> RouterModelInfo: