diff --git a/litellm/proxy/guardrails/guardrail_hooks/javelin/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/javelin/__init__.py new file mode 100644 index 00000000000..7dcb190a844 --- /dev/null +++ b/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, +} diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 3f9f7cec0e0..03347cd1cdd 100644 --- a/litellm/types/guardrails.py +++ b/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")