fix: pattern for guardrail registration

litellm/proxy/guardrails/guardrail_hooks/javelin/__init__.py

@@ -0,0 +1,43 @@
+from typing import TYPE_CHECKING
+
+from litellm.types.guardrails import SupportedGuardrailIntegrations
+
+from .javelin import JavelinGuardrail
+
+if TYPE_CHECKING:
+    from litellm.types.guardrails import Guardrail, LitellmParams
+
+
+def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"):
+    import litellm
+
+    if litellm_params.guard_name is None:
+        raise Exception(
+            "JavelinGuardrailException - Please pass the Javelin guard name via 'litellm_params::guard_name'"
+        )
+
+    _javelin_callback = JavelinGuardrail(
+        api_base=litellm_params.api_base,
+        api_key=litellm_params.api_key,
+        guardrail_name=guardrail.get("guardrail_name", ""),
+        javelin_guard_name=litellm_params.guard_name,
+        event_hook=litellm_params.mode,
+        default_on=litellm_params.default_on,
+        api_version=litellm_params.api_version,
+        config=litellm_params.config,
+        metadata=litellm_params.metadata,
+        application=litellm_params.application,
+    )
+    litellm.logging_callback_manager.add_litellm_callback(_javelin_callback)
+
+    return _javelin_callback
+
+
+guardrail_initializer_registry = {
+    SupportedGuardrailIntegrations.JAVELIN.value: initialize_guardrail,
+}
+
+
+guardrail_class_registry = {
+    SupportedGuardrailIntegrations.JAVELIN.value: JavelinGuardrail,
+}

litellm/types/guardrails.py

@@ -38,6 +38,7 @@ class SupportedGuardrailIntegrations(Enum):
     OPENAI_MODERATION = "openai_moderation"
     NOMA = "noma"
     TOOL_PERMISSION = "tool_permission"
+    JAVELIN = "javelin"
 
 
 class Role(Enum):
@@ -390,6 +391,26 @@ class ToolPermissionGuardrailConfigModel(BaseModel):
     )
 
 
+class JavelinGuardrailConfigModel(BaseModel):
+    """Configuration parameters for the Javelin guardrail"""
+
+    guard_name: Optional[str] = Field(
+        default=None, description="Name of the Javelin guard to use"
+    )
+    api_version: Optional[str] = Field(
+        default="v1", description="API version for Javelin service"
+    )
+    metadata: Optional[Dict] = Field(
+        default=None, description="Additional metadata to send with requests"
+    )
+    application: Optional[str] = Field(
+        default=None, description="Application name for Javelin service"
+    )
+    config: Optional[Dict] = Field(
+        default=None, description="Additional configuration for the guardrail"
+    )
+
+
 class BaseLitellmParams(BaseModel):  # works for new and patch update guardrails
     api_key: Optional[str] = Field(
         default=None, description="API key for the guardrail service"
@@ -479,6 +500,7 @@ class LitellmParams(
     PillarGuardrailConfigModel,
     NomaGuardrailConfigModel,
     ToolPermissionGuardrailConfigModel,
+    JavelinGuardrailConfigModel,
     BaseLitellmParams,
 ):
     guardrail: str = Field(description="The type of guardrail integration to use")