chore: initial snapshot for gitea/github upload
This commit is contained in:
Your Name
@@ -0,0 +1,38 @@
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from litellm.types.guardrails import SupportedGuardrailIntegrations
|
||||
|
||||
from .hiddenlayer import HiddenlayerGuardrail
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.types.guardrails import Guardrail, LitellmParams
|
||||
|
||||
|
||||
def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"):
|
||||
import litellm
|
||||
|
||||
api_id = litellm_params.api_id if hasattr(litellm_params, "api_id") else None
|
||||
auth_url = litellm_params.auth_url if hasattr(litellm_params, "auth_url") else None
|
||||
|
||||
_hiddenlayer_callback = HiddenlayerGuardrail(
|
||||
api_base=litellm_params.api_base,
|
||||
api_id=api_id,
|
||||
api_key=litellm_params.api_key,
|
||||
auth_url=auth_url,
|
||||
guardrail_name=guardrail.get("guardrail_name", ""),
|
||||
event_hook=litellm_params.mode,
|
||||
default_on=litellm_params.default_on,
|
||||
)
|
||||
|
||||
litellm.logging_callback_manager.add_litellm_callback(_hiddenlayer_callback)
|
||||
return _hiddenlayer_callback
|
||||
|
||||
|
||||
guardrail_initializer_registry = {
|
||||
SupportedGuardrailIntegrations.HIDDENLAYER.value: initialize_guardrail,
|
||||
}
|
||||
|
||||
|
||||
guardrail_class_registry = {
|
||||
SupportedGuardrailIntegrations.HIDDENLAYER.value: HiddenlayerGuardrail,
|
||||
}
|
||||
@@ -0,0 +1,259 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from typing import TYPE_CHECKING, Any, Literal, Optional, Type
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import requests
|
||||
from fastapi import HTTPException
|
||||
from httpx import HTTPStatusError
|
||||
from requests.auth import HTTPBasicAuth
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.integrations.custom_guardrail import (
|
||||
CustomGuardrail,
|
||||
log_guardrail_information,
|
||||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
get_async_httpx_client,
|
||||
httpxSpecialProvider,
|
||||
)
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.hiddenlayer import (
|
||||
HiddenlayerAction,
|
||||
HiddenlayerMessages,
|
||||
)
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
|
||||
|
||||
|
||||
def is_saas(host: str) -> bool:
|
||||
"""Checks whether the connection is to the SaaS platform"""
|
||||
|
||||
o = urlparse(host)
|
||||
|
||||
if o.hostname and o.hostname.endswith("hiddenlayer.ai"):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def _get_jwt(auth_url, api_id, api_key):
|
||||
token_url = f"{auth_url}/oauth2/token?grant_type=client_credentials"
|
||||
|
||||
resp = requests.post(token_url, auth=HTTPBasicAuth(api_id, api_key))
|
||||
|
||||
if not resp.ok:
|
||||
raise RuntimeError(
|
||||
f"Unable to get authentication credentials for the HiddenLayer API: {resp.status_code}: {resp.text}"
|
||||
)
|
||||
|
||||
if "access_token" not in resp.json():
|
||||
raise RuntimeError(
|
||||
f"Unable to get authentication credentials for the HiddenLayer API - invalid response: {resp.json()}"
|
||||
)
|
||||
|
||||
return resp.json()["access_token"]
|
||||
|
||||
|
||||
class HiddenlayerGuardrail(CustomGuardrail):
|
||||
"""Custom guardrail wrapper for HiddenLayer's safety checks."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_id: Optional[str] = None,
|
||||
api_key: Optional[str] = None,
|
||||
api_base: Optional[str] = None,
|
||||
auth_url: Optional[str] = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
self.hiddenlayer_client_id = api_id or os.getenv("HIDDENLAYER_CLIENT_ID")
|
||||
self.hiddenlayer_client_secret = api_key or os.getenv(
|
||||
"HIDDENLAYER_CLIENT_SECRET"
|
||||
)
|
||||
self.api_base = (
|
||||
api_base
|
||||
or os.getenv("HIDDENLAYER_API_BASE")
|
||||
or "https://api.hiddenlayer.ai"
|
||||
)
|
||||
self.jwt_token = None
|
||||
|
||||
auth_url = (
|
||||
auth_url
|
||||
or os.getenv("HIDDENLAYER_AUTH_URL")
|
||||
or "https://auth.hiddenlayer.ai"
|
||||
)
|
||||
|
||||
if is_saas(self.api_base):
|
||||
if not self.hiddenlayer_client_id:
|
||||
raise RuntimeError(
|
||||
"`api_id` cannot be None when using the SaaS version of HiddenLayer."
|
||||
)
|
||||
|
||||
if not self.hiddenlayer_client_secret:
|
||||
raise RuntimeError(
|
||||
"`api_key` cannot be None when using the SaaS version of HiddenLayer."
|
||||
)
|
||||
|
||||
self.jwt_token = _get_jwt(
|
||||
auth_url=auth_url,
|
||||
api_id=self.hiddenlayer_client_id,
|
||||
api_key=self.hiddenlayer_client_secret,
|
||||
)
|
||||
self.refresh_jwt_func = lambda: _get_jwt(
|
||||
auth_url=auth_url,
|
||||
api_id=self.hiddenlayer_client_id,
|
||||
api_key=self.hiddenlayer_client_secret,
|
||||
)
|
||||
|
||||
self._http_client = get_async_httpx_client(
|
||||
llm_provider=httpxSpecialProvider.GuardrailCallback
|
||||
)
|
||||
super().__init__(**kwargs)
|
||||
|
||||
@log_guardrail_information
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict,
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: Optional["LiteLLMLoggingObj"] = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
"""Validate (and optionally redact) text via HiddenLayer before/after LLM calls."""
|
||||
|
||||
# The model in the request and the response can be inconsistent
|
||||
# I.e request can specify gpt-4o-mini but the response from the server will be
|
||||
# gpt-4o-mini-2025-11-01. We need the model to be consistent so that inferences
|
||||
# will be grouped correctly on the Hiddenlayer side
|
||||
model_name = (
|
||||
logging_obj.model if logging_obj and logging_obj.model else "unknown"
|
||||
)
|
||||
hl_request_metadata = {"model": model_name}
|
||||
|
||||
# We need the hiddenlayer project id and requester id on both the input and output
|
||||
# Since headers aren't available on the response back from the model, we get them
|
||||
# from the logging object. It ends up working out that on the request, we parse the
|
||||
# hiddenlayer params from the raw request and then retrieve those same headers
|
||||
# from the logger object on the response from the model.
|
||||
headers = request_data.get("proxy_server_request", {}).get("headers", {})
|
||||
if not headers and logging_obj and logging_obj.model_call_details:
|
||||
headers = (
|
||||
logging_obj.model_call_details.get("litellm_params", {})
|
||||
.get("metadata", {})
|
||||
.get("headers", {})
|
||||
)
|
||||
|
||||
hl_request_metadata["requester_id"] = (
|
||||
headers.get("hl-requester-id") or "LiteLLM"
|
||||
)
|
||||
project_id = headers.get("hl-project-id")
|
||||
|
||||
if scan_params := inputs.get("structured_messages"):
|
||||
# Convert AllMessageValues to simple dict format for HiddenLayer API
|
||||
messages = [
|
||||
{"role": msg.get("role", "user"), "content": msg.get("content", "")}
|
||||
for msg in scan_params
|
||||
if isinstance(msg, dict)
|
||||
]
|
||||
result = await self._call_hiddenlayer(
|
||||
project_id, hl_request_metadata, {"messages": messages}, input_type
|
||||
)
|
||||
elif text := inputs.get("texts"):
|
||||
result = await self._call_hiddenlayer(
|
||||
project_id,
|
||||
hl_request_metadata,
|
||||
{"messages": [{"role": "user", "content": text[-1]}]},
|
||||
input_type,
|
||||
)
|
||||
else:
|
||||
result = {}
|
||||
|
||||
if result.get("evaluation", {}).get("action") == HiddenlayerAction.BLOCK:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
"error": "Violated guardrail policy",
|
||||
"hiddenlayer_guardrail_response": HiddenlayerMessages.BLOCK_MESSAGE,
|
||||
},
|
||||
)
|
||||
|
||||
if result.get("evaluation", {}).get("action") == HiddenlayerAction.REDACT:
|
||||
modified_data = result.get("modified_data", {})
|
||||
if modified_data.get("input") and input_type == "request":
|
||||
inputs["texts"] = [modified_data["input"]["messages"][-1]["content"]]
|
||||
inputs["structured_messages"] = modified_data["input"]["messages"]
|
||||
|
||||
if modified_data.get("output") and input_type == "response":
|
||||
inputs["texts"] = [modified_data["output"]["messages"][-1]["content"]]
|
||||
|
||||
return inputs
|
||||
|
||||
async def _call_hiddenlayer(
|
||||
self,
|
||||
project_id: str | None,
|
||||
metadata: dict[str, str],
|
||||
payload: dict[str, Any],
|
||||
input_type: Literal["request", "response"],
|
||||
) -> dict[str, Any]:
|
||||
data: dict[str, Any] = {"metadata": metadata}
|
||||
|
||||
if input_type == "request":
|
||||
data["input"] = payload
|
||||
else:
|
||||
data["output"] = payload
|
||||
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
if project_id:
|
||||
headers["HL-Project-Id"] = project_id
|
||||
|
||||
if self.jwt_token:
|
||||
headers["Authorization"] = f"Bearer {self.jwt_token}"
|
||||
|
||||
try:
|
||||
response = await self._http_client.post(
|
||||
f"{self.api_base}/detection/v1/interactions",
|
||||
json=data,
|
||||
headers=headers,
|
||||
)
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
|
||||
verbose_proxy_logger.debug(f"Hiddenlayer reponse: {result}")
|
||||
|
||||
return result
|
||||
except HTTPStatusError as e:
|
||||
# Try the request again by refreshing the jwt if we get 401
|
||||
# since the Hiddenlayer jwt timeout is an hour and this is
|
||||
# a long lived session application
|
||||
if e.response.status_code == 401 and self.jwt_token is not None:
|
||||
verbose_proxy_logger.debug(
|
||||
"Unable to authenticate to Hiddenlayer, JWT token is invalid or expired, trying to refresh the token."
|
||||
)
|
||||
self.jwt_token = self.refresh_jwt_func()
|
||||
headers["Authorization"] = f"Bearer {self.jwt_token}"
|
||||
response = await self._http_client.post(
|
||||
f"{self.api_base}/detection/v1/interactions",
|
||||
json=data,
|
||||
headers=headers,
|
||||
)
|
||||
else:
|
||||
raise e
|
||||
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
|
||||
verbose_proxy_logger.debug(f"Hiddenlayer reponse: {result}")
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def get_config_model() -> Optional[Type["GuardrailConfigModel"]]:
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.hiddenlayer import (
|
||||
HiddenlayerGuardrailConfigModel,
|
||||
)
|
||||
|
||||
return HiddenlayerGuardrailConfigModel
|
||||
Reference in New Issue
Block a user