diff --git a/docs/my-website/docs/proxy/guardrails/ibm_guardrails.md b/docs/my-website/docs/proxy/guardrails/ibm_guardrails.md index 0c13d2dcea9..43ba6622078 100644 --- a/docs/my-website/docs/proxy/guardrails/ibm_guardrails.md +++ b/docs/my-website/docs/proxy/guardrails/ibm_guardrails.md @@ -95,6 +95,7 @@ curl -i http://localhost:4000/v1/chat/completions \ These go under `optional_params`: - `detector_params` - dict - Parameters to pass to your detector +- `extra_headers` - dict - Additional headers to inject into requests to IBM Guardrails, as a key-value dict. - `score_threshold` - float - Only count detections above this score (0.0 to 1.0) - `block_on_detection` - bool - Block the request when violations found. Default: `true` diff --git a/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/__init__.py index b22b90d3679..e397d8098a1 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/__init__.py @@ -1,6 +1,7 @@ from typing import TYPE_CHECKING from litellm.types.guardrails import SupportedGuardrailIntegrations +from litellm.types.proxy.guardrails.guardrail_hooks.ibm import IBMDetectorOptionalParams from .ibm_detector import IBMGuardrailDetector @@ -22,11 +23,17 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" if not guardrail_name: raise ValueError("IBM Guardrails: guardrail_name is required") - # Get optional params - detector_params = getattr(litellm_params, "detector_params", {}) - score_threshold = getattr(litellm_params, "score_threshold", None) - block_on_detection = getattr(litellm_params, "block_on_detection", True) verify_ssl = getattr(litellm_params, "verify_ssl", True) + + # Get optional params + optional_params = getattr(litellm_params, "optional_params", IBMDetectorOptionalParams()) + detector_params = getattr(optional_params, "detector_params", {}) + extra_headers = getattr(optional_params, "extra_headers", {}) + score_threshold = getattr(optional_params, "score_threshold", None) + block_on_detection = getattr(optional_params, "block_on_detection", True) + + + is_detector_server = litellm_params.is_detector_server if is_detector_server is None: is_detector_server = True @@ -38,6 +45,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" detector_id=litellm_params.detector_id, is_detector_server=is_detector_server, detector_params=detector_params, + extra_headers=extra_headers, score_threshold=score_threshold, block_on_detection=block_on_detection, verify_ssl=verify_ssl, diff --git a/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py b/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py index fefd39b7214..55fa17c21e7 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py @@ -39,6 +39,7 @@ class IBMGuardrailDetector(CustomGuardrail): detector_id: Optional[str] = None, is_detector_server: bool = True, detector_params: Optional[Dict[str, Any]] = None, + extra_headers: Optional[Dict[str, str]] = None, score_threshold: Optional[float] = None, block_on_detection: bool = True, verify_ssl: bool = True, @@ -70,6 +71,7 @@ class IBMGuardrailDetector(CustomGuardrail): self.is_detector_server = is_detector_server self.detector_params = detector_params or {} + self.extra_headers = extra_headers or {} self.score_threshold = score_threshold self.block_on_detection = block_on_detection self.verify_ssl = verify_ssl @@ -129,6 +131,10 @@ class IBMGuardrailDetector(CustomGuardrail): "detector-id": self.detector_id, } + # Add any extra headers to the request + for header, value in self.extra_headers.items(): + headers[header] = value + verbose_proxy_logger.debug( "IBM Detector Server request to %s with payload: %s", self.api_url, @@ -217,6 +223,10 @@ class IBMGuardrailDetector(CustomGuardrail): "content-type": "application/json", } + # Add any extra headers to the request + for header, value in self.extra_headers.items(): + headers[header] = value + verbose_proxy_logger.debug( "IBM Orchestrator request to %s with payload: %s", self.api_url, diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/ibm/ibm_detector.py b/litellm/types/proxy/guardrails/guardrail_hooks/ibm/ibm_detector.py index e2f2febb7a7..e30f8c938ae 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/ibm/ibm_detector.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/ibm/ibm_detector.py @@ -57,7 +57,12 @@ class IBMDetectorOptionalParams(BaseModel): detector_params: Optional[Dict[str, Any]] = Field( default_factory=lambda: {}, - description="Dictionary of arguments to pass to the detector. Can be set per-request or hard-coded at guardrail config time.", + description="Dictionary of arguments to pass to the detector.", + ) + + extra_headers: Optional[Dict[str, Any]] = Field( + default_factory=lambda: {}, + description="Dictionary of extra headers to pass to the detector.", ) score_threshold: Optional[float] = Field(