Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/guardrails/guardrail

1from __future__ import annotations

3import os

4from collections.abc import Mapping, Sequence

5from typing import TYPE_CHECKING, Final, Literal, Protocol

6from urllib.parse import urlparse

7from uuid import uuid4

9import httpx

10import requests

11from fastapi import HTTPException

12from httpx import HTTPStatusError

13from requests.auth import HTTPBasicAuth

14from typing_extensions import ReadOnly, TypedDict, Unpack

16from litellm._logging import verbose_proxy_logger

17from litellm.integrations.custom_guardrail import (

18 CustomGuardrail,

19 log_guardrail_information,

20)

21from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj

22from litellm.llms.custom_httpx.http_handler import (

23 get_async_httpx_client,

24 httpxSpecialProvider,

25)

26from litellm.types.guardrails import GuardrailEventHooks

27from litellm.types.llms.openai import ChatCompletionToolCallChunk

28from litellm.types.proxy.guardrails.guardrail_hooks.hiddenlayer import (

29 HiddenlayerAction,

30 HiddenlayerMessages,

31)

32from litellm.types.utils import ChatCompletionMessageToolCall, GenericGuardrailAPIInputs

34if TYPE_CHECKING: 34 ↛ 35line 34 didn't jump to line 35 because the condition on line 34 was never true

35 from pydantic import BaseModel

37 from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel

40_AUTH_TIMEOUT_SECONDS: Final[float] = 30.0

43class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object):

44 """Base-class constructor options carried by this guardrail's forwarded keyword arguments."""

46 guardrail_name: ReadOnly[str | None]

47 supported_event_hooks: list[GuardrailEventHooks] | None

50class _HiddenlayerEvaluation(TypedDict, total=False):

51 action: ReadOnly[str]

52 threat_level: ReadOnly[str]

55class _HiddenlayerAnalysisEntry(TypedDict, total=False):

56 name: ReadOnly[str]

57 detected: ReadOnly[bool]

60class _HiddenlayerModifiedMessage(TypedDict):

61 content: ReadOnly[str | list[Mapping[str, str]]]

64class _HiddenlayerModifiedSide(TypedDict):

65 messages: ReadOnly[list[_HiddenlayerModifiedMessage]]

68class _HiddenlayerResponse(TypedDict, total=False):

69 evaluation: ReadOnly[_HiddenlayerEvaluation]

70 analysis: ReadOnly[Sequence[_HiddenlayerAnalysisEntry]]

71 modified_data: ReadOnly[Mapping[str, _HiddenlayerModifiedSide]]

74class _ProxyServerRequest(TypedDict, total=False):

75 headers: ReadOnly[dict[str, str]]

78class _HiddenlayerRequestData(TypedDict, total=False):

79 proxy_server_request: ReadOnly[_ProxyServerRequest]

82class _LoggedCallMetadata(TypedDict, total=False):

83 headers: ReadOnly[dict[str, str]]

86class _LoggedCallLitellmParams(TypedDict, total=False):

87 metadata: ReadOnly[_LoggedCallMetadata]

90class _HiddenlayerOutputMessage(TypedDict, total=False):

91 content: ReadOnly[str | list[Mapping[str, str]]]

94class _HiddenlayerChoiceMessage(TypedDict, total=False):

95 content: ReadOnly[str]

98class _HiddenlayerChoice(TypedDict, total=False):

99 message: ReadOnly[_HiddenlayerChoiceMessage]

100

101

102class _HiddenlayerV2Output(TypedDict, total=False):

103 messages: ReadOnly[Sequence[_HiddenlayerOutputMessage]]

104 choices: ReadOnly[Sequence[_HiddenlayerChoice]]

105

106

107class _HiddenlayerV2OutputView(TypedDict):

108 """Typed read of the untyped JSON body returned by the HiddenLayer detection endpoints."""

109

110 evaluation: ReadOnly[_HiddenlayerV2Output]

111

112

113_HiddenlayerV2Payload = Mapping[str, object] | list[ChatCompletionToolCallChunk] | list[ChatCompletionMessageToolCall]

114

115

116class _LoggedCallDetails(Protocol):

117 """Logging object view that exposes its untyped call details with the shape this guardrail reads."""

118

119 @property

120 def model_call_details(self) -> Mapping[str, _LoggedCallLitellmParams]: ... 120 ↛ exitline 120 didn't return from function 'model_call_details' because

121

122

123class _TokenPayloadSource(Protocol):

124 """Response view that decodes the HiddenLayer OAuth token body as a string mapping."""

125

126 def json(self) -> Mapping[str, str]: ... 126 ↛ exitline 126 didn't return from function 'json' because

127

128

129class _InteractionPayloadSource(Protocol):

130 """Response view that decodes the HiddenLayer v1 interaction body with the shape this guardrail reads."""

131

132 def json(self) -> _HiddenlayerResponse: ... 132 ↛ exitline 132 didn't return from function 'json' because

133

134

135def _interaction_body(response: _InteractionPayloadSource) -> _HiddenlayerResponse:

136 return response.json()

137

138

139def _proxy_server_request(request_data: _HiddenlayerRequestData) -> _ProxyServerRequest | None:

140 return request_data.get("proxy_server_request")

141

142

143def _proxy_request_headers(request_data: _HiddenlayerRequestData) -> dict[str, str]:

144 return request_data.get("proxy_server_request", {}).get("headers", {})

145

146

147def _logged_request_headers(logging_obj: _LoggedCallDetails) -> dict[str, str]:

148 return logging_obj.model_call_details.get("litellm_params", {}).get("metadata", {}).get("headers", {})

149

150

151def _token_payload(response: _TokenPayloadSource) -> Mapping[str, str]:

152 return response.json()

153

154

155def _header_value(headers: Mapping[str, str], key: str, default: str) -> str:

156 return headers.get(key, default)

157

158

159def _is_image_part(item: object) -> bool:

160 """Whether a structured-message content part carries an image rather than text."""

161

162 if not isinstance(item, Mapping):

163 return False

164

165 part: Final[Mapping[object, object]] = item

166 return part.get("type") == "image_url"

167

168

169def _scannable_text(content: object) -> str:

170 """Flatten a structured message's content into the single string the v1 detection endpoint takes.

171

172 Image parts are dropped: the endpoint accepts one string, so an image would only reach it as

173 its stringified source (a base64 blob or a URL), which is not text the scanner can evaluate.

174 """

175

176 if not isinstance(content, list):

177 return str(content or "")

178

179 parts: Final[Sequence[object]] = content

180 text_parts: Final = [item for item in parts if not _is_image_part(item)] # mutable-ok: sent as a list repr

181 return str(text_parts or "")

182

183

184def is_saas(host: str) -> bool:

185 """Checks whether the connection is to the SaaS platform"""

186

187 o: Final = urlparse(host)

188

189 if o.hostname and o.hostname.endswith("hiddenlayer.ai"):

190 return True

191

192 return False

193

194

195def _get_jwt(auth_url, api_id, api_key, timeout: float = _AUTH_TIMEOUT_SECONDS) -> str:

196 token_url: Final = f"{auth_url}/oauth2/token?grant_type=client_credentials"

197

198 resp: Final = requests.post(token_url, auth=HTTPBasicAuth(api_id, api_key), timeout=timeout)

199

200 if not resp.ok:

201 raise RuntimeError(

202 f"Unable to get authentication credentials for the HiddenLayer API: {resp.status_code}: {resp.text}"

203 )

204

205 if "access_token" not in resp.json():

206 raise RuntimeError(

207 f"Unable to get authentication credentials for the HiddenLayer API - invalid response: {resp.json()}"

208 )

209

210 return _token_payload(resp)["access_token"]

211

212

213class HiddenlayerGuardrail(CustomGuardrail):

214 """Custom guardrail wrapper for HiddenLayer's safety checks."""

215

216 @classmethod

217 def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]:

218 return [

219 GuardrailEventHooks.pre_call,

220 GuardrailEventHooks.post_call,

221 ]

222

223 def __init__(

224 self,

225 api_id: str | None = None,

226 api_key: str | None = None,

227 api_base: str | None = None,

228 auth_url: str | None = None,

229 **kwargs: Unpack[_CustomGuardrailOptions],

230 ) -> None:

231 kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks()))

232 self.hiddenlayer_client_id = api_id or os.getenv("HIDDENLAYER_CLIENT_ID")

233 self.hiddenlayer_client_secret = api_key or os.getenv("HIDDENLAYER_CLIENT_SECRET")

234 self.api_base = api_base or os.getenv("HIDDENLAYER_API_BASE") or "https://api.hiddenlayer.ai"

235 self.jwt_token = None

236

237 auth_url = auth_url or os.getenv("HIDDENLAYER_AUTH_URL") or "https://auth.hiddenlayer.ai"

238

239 if is_saas(self.api_base):

240 if not self.hiddenlayer_client_id:

241 raise RuntimeError("`api_id` cannot be None when using the SaaS version of HiddenLayer.")

242

243 if not self.hiddenlayer_client_secret:

244 raise RuntimeError("`api_key` cannot be None when using the SaaS version of HiddenLayer.")

245

246 self.jwt_token = _get_jwt(

247 auth_url=auth_url,

248 api_id=self.hiddenlayer_client_id,

249 api_key=self.hiddenlayer_client_secret,

250 )

251 self.refresh_jwt_func = lambda: _get_jwt(

252 auth_url=auth_url,

253 api_id=self.hiddenlayer_client_id,

254 api_key=self.hiddenlayer_client_secret,

255 )

256

257 self._http_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)

258 super().__init__(**kwargs)

259

260 @log_guardrail_information

261 async def apply_guardrail(

262 self,

263 inputs: GenericGuardrailAPIInputs,

264 request_data: dict,

265 input_type: Literal["request", "response"],

266 logging_obj: LiteLLMLoggingObj | None = None,

267 ) -> GenericGuardrailAPIInputs:

268 """Validate (and optionally redact) text via HiddenLayer before/after LLM calls."""

269

270 # The model in the request and the response can be inconsistent

271 # I.e request can specify gpt-4o-mini but the response from the server will be

272 # gpt-4o-mini-2025-11-01. We need the model to be consistent so that inferences

273 # will be grouped correctly on the Hiddenlayer side

274 model_name: Final = logging_obj.model if logging_obj and logging_obj.model else "unknown"

275 hl_request_metadata: Final = {"model": model_name}

276

277 # We need the hiddenlayer project id and requester id on both the input and output

278 # Since headers aren't available on the response back from the model, we get them

279 # from the logging object. It ends up working out that on the request, we parse the

280 # hiddenlayer params from the raw request and then retrieve those same headers

281 # from the logger object on the response from the model.

282 headers = _proxy_request_headers(request_data)

283 if not headers and logging_obj and logging_obj.model_call_details:

284 headers = _logged_request_headers(logging_obj)

285

286 hl_request_metadata["requester_id"] = headers.get("hl-requester-id") or "LiteLLM"

287 project_id: Final = headers.get("hl-project-id")

288

289 if scan_params := inputs.get("structured_messages"):

290 last_msg: Final = scan_params[-1]

291 result: _HiddenlayerResponse = await self._call_hiddenlayer(

292 project_id,

293 hl_request_metadata,

294 {

295 "messages": [

296 {

297 "role": last_msg.get("role", "user"),

298 "content": _scannable_text(last_msg.get("content")),

299 }

300 ]

301 },

302 input_type,

303 )

304 elif text := inputs.get("texts"):

305 result = await self._call_hiddenlayer(

306 project_id,

307 hl_request_metadata,

308 {"messages": [{"role": "user", "content": text[-1]}]},

309 input_type,

310 )

311 else:

312 result = {}

313

314 if result.get("evaluation", {}).get("action") == HiddenlayerAction.BLOCK:

315 detected_reasons: Final = [

316 entry.get("name", "unknown") for entry in result.get("analysis", []) if entry.get("detected")

317 ]

318 threat_level: Final = result.get("evaluation", {}).get("threat_level")

319 raise HTTPException(

320 status_code=400,

321 detail={

322 "error": "Violated guardrail policy",

323 "hiddenlayer_guardrail_response": HiddenlayerMessages.BLOCK_MESSAGE.value,

324 "block_reasons": detected_reasons,

325 "threat_level": threat_level,

326 },

327 )

328

329 if result.get("evaluation", {}).get("action") == HiddenlayerAction.REDACT:

330 modified_data: Final = result.get("modified_data", {})

331 if modified_data.get("input") and input_type == "request":

332 last_content = modified_data["input"]["messages"][-1]["content"]

333 if isinstance(last_content, list):

334 texts = [

335 item["text"] for item in last_content if isinstance(item, dict) and item.get("type") == "text"

336 ]

337 inputs["texts"] = texts if texts else [""]

338 else:

339 inputs["texts"] = [last_content]

340 inputs["structured_messages"] = modified_data["input"]["messages"]

341

342 if modified_data.get("output") and input_type == "response":

343 last_content = modified_data["output"]["messages"][-1]["content"]

344 if isinstance(last_content, list):

345 texts = [

346 item["text"] for item in last_content if isinstance(item, dict) and item.get("type") == "text"

347 ]

348 inputs["texts"] = texts if texts else [""]

349 else:

350 inputs["texts"] = [last_content]

351

352 return inputs

353

354 async def _call_hiddenlayer(

355 self,

356 project_id: str | None,

357 metadata: Mapping[str, str],

358 payload: Mapping[str, Sequence[Mapping[str, str]]],

359 input_type: Literal["request", "response"],

360 ) -> _HiddenlayerResponse:

361 data: Final[dict[str, object]] = {"metadata": metadata}

362

363 if input_type == "request":

364 data["input"] = payload

365 else:

366 data["output"] = payload

367

368 headers: Final = {

369 "Content-Type": "application/json",

370 "hl-runtime-edge-provider": "litellm",

371 "hl-runtime-edge-provider-version": "1",

372 }

373

374 if project_id:

375 headers["HL-Project-Id"] = project_id

376

377 if self.jwt_token:

378 headers["Authorization"] = f"Bearer {self.jwt_token}"

379

380 try:

381 response = await self._http_client.post(

382 f"{self.api_base}/detection/v1/interactions",

383 json=data,

384 headers=headers,

385 )

386 response.raise_for_status()

387 result: _HiddenlayerResponse = _interaction_body(response)

388

389 verbose_proxy_logger.debug("Hiddenlayer reponse: %s", result)

390

391 return result

392 except HTTPStatusError as e:

393 # Try the request again by refreshing the jwt if we get 401

394 # since the Hiddenlayer jwt timeout is an hour and this is

395 # a long lived session application

396 if e.response.status_code == 401 and self.jwt_token is not None:

397 verbose_proxy_logger.debug(

398 "Unable to authenticate to Hiddenlayer, JWT token is invalid or expired, trying to refresh the token."

399 )

400 self.jwt_token = self.refresh_jwt_func()

401 headers["Authorization"] = f"Bearer {self.jwt_token}"

402 response = await self._http_client.post(

403 f"{self.api_base}/detection/v1/interactions",

404 json=data,

405 headers=headers,

406 )

407 else:

408 raise e

409

410 response.raise_for_status()

411 result = _interaction_body(response)

412

413 verbose_proxy_logger.debug("Hiddenlayer reponse: %s", result)

414 return result

415

416 @staticmethod

417 def get_config_model() -> type[GuardrailConfigModel[BaseModel]] | None:

418 from litellm.types.proxy.guardrails.guardrail_hooks.hiddenlayer import (

419 HiddenlayerGuardrailConfigModel,

420 )

421

422 return HiddenlayerGuardrailConfigModel

423

424

425class HiddenlayerGuardrailV2(CustomGuardrail):

426 """Custom guardrail wrapper for HiddenLayer's safety checks."""

427

428 def __init__(

429 self,

430 api_id: str | None = None,

431 api_key: str | None = None,

432 api_base: str | None = None,

433 auth_url: str | None = None,

434 **kwargs: Unpack[_CustomGuardrailOptions],

435 ) -> None:

436 self.hiddenlayer_client_id = api_id or os.getenv("HIDDENLAYER_CLIENT_ID")

437 self.hiddenlayer_client_secret = api_key or os.getenv("HIDDENLAYER_CLIENT_SECRET")

438 self.api_base = api_base or os.getenv("HIDDENLAYER_API_BASE") or "https://api.hiddenlayer.ai"

439 self.jwt_token = None

440

441 auth_url = auth_url or os.getenv("HIDDENLAYER_AUTH_URL") or "https://auth.hiddenlayer.ai"

442

443 if is_saas(self.api_base):

444 if not self.hiddenlayer_client_id:

445 raise RuntimeError("`api_id` cannot be None when using the SaaS version of HiddenLayer.")

446

447 if not self.hiddenlayer_client_secret:

448 raise RuntimeError("`api_key` cannot be None when using the SaaS version of HiddenLayer.")

449

450 self.jwt_token = _get_jwt(

451 auth_url=auth_url,

452 api_id=self.hiddenlayer_client_id,

453 api_key=self.hiddenlayer_client_secret,

454 )

455 self.refresh_jwt_func = lambda: _get_jwt(

456 auth_url=auth_url,

457 api_id=self.hiddenlayer_client_id,

458 api_key=self.hiddenlayer_client_secret,

459 )

460

461 self._http_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)

462 super().__init__(**kwargs)

463

464 @log_guardrail_information

465 async def apply_guardrail(

466 self,

467 inputs: GenericGuardrailAPIInputs,

468 request_data: dict,

469 input_type: Literal["request", "response"],

470 logging_obj: LiteLLMLoggingObj | None = None,

471 ) -> GenericGuardrailAPIInputs:

472 """Validate (and optionally redact) text via HiddenLayer before/after LLM calls."""

473

474 # We need the hiddenlayer project id and requester id on both the input and output

475 # Since headers aren't available on the response back from the model, we get them

476 # from the logging object. It ends up working out that on the request, we parse the

477 # hiddenlayer params from the raw request and then retrieve those same headers

478 # from the logger object on the response from the model.

479 headers = _proxy_request_headers(request_data)

480 if not headers and logging_obj and logging_obj.model_call_details:

481 headers = _logged_request_headers(logging_obj)

482

483 # put our roundtrip id in the header to the model so we get it on the way back from the model

484 if "hl-roundtrip-id" not in headers:

485 proxy_req: Final = _proxy_server_request(request_data)

486 if proxy_req is not None and "headers" in proxy_req:

487 proxy_req["headers"]["hl-roundtrip-id"] = str(uuid4())

488 headers["hl-roundtrip-id"] = proxy_req["headers"]["hl-roundtrip-id"]

489

490 hl_headers: Final = {h.lower(): v for h, v in headers.items() if h.lower().startswith("hl-")}

491

492 if "hl-requester-id" not in hl_headers:

493 hl_headers["hl-requester-id"] = "LiteLLM"

494

495 payload: _HiddenlayerV2Payload

496 if input_type == "request":

497 payload = {

498 "messages": inputs.get("structured_messages"),

499 "model": inputs.get("model"),

500 "tools": inputs.get("tools"),

501 }

502 else:

503 if inputs.get("texts"):

504 payload = {

505 "choices": [

506 {

507 "index": 0,

508 "message": {

509 "role": "assistant",

510 "content": (inputs["texts"][0] if inputs.get("texts") else ""),

511 },

512 "finish_reason": "stop",

513 }

514 ]

515 }

516 elif tool_calls := inputs.get("tool_calls"):

517 payload = tool_calls

518 else:

519 payload = {}

520

521 response: Final = await self._call_hiddenlayer(payload, input_type, hl_headers)

522 output: Final = response.json()

523 output_view: Final[_HiddenlayerV2OutputView] = {"evaluation": output}

524 evaluated_output: Final = output_view["evaluation"]

525

526 if _header_value(response.headers, "hl-runtime-action", "").lower() == "block":

527 raise HTTPException(

528 status_code=400,

529 detail={

530 "error": "Violated guardrail policy",

531 "hiddenlayer_guardrail_response": HiddenlayerMessages.BLOCK_MESSAGE.value,

532 },

533 )

534

535 new_texts: Final[list[str]] = []

536 if input_type == "request":

537 inputs["structured_messages"] = output

538

539 modified_messages: Final[Sequence[_HiddenlayerOutputMessage]] = evaluated_output.get("messages", [])

540 for message in modified_messages:

541 content = message.get("content", "")

542 if isinstance(content, list):

543 text_parts = [

544 item["text"] for item in content if isinstance(item, dict) and item.get("type") == "text"

545 ]

546 if text_parts:

547 new_texts.append(" ".join(text_parts))

548 elif content:

549 new_texts.append(content)

550

551 inputs["texts"] = new_texts

552

553 elif input_type == "response" and inputs.get("texts"):

554 redacted_choices: Final[Sequence[_HiddenlayerChoice]] = evaluated_output.get("choices", [{}])

555 inputs["texts"] = [redacted_choices[-1].get("message", {}).get("content", "")]

556 elif input_type == "response" and inputs.get("tool_calls"):

557 inputs["tool_calls"] = output

558

559 return inputs

560

561 async def _call_hiddenlayer(

562 self,

563 payload: _HiddenlayerV2Payload,

564 input_type: Literal["request", "response"],

565 hl_headers: Mapping[str, str],

566 ) -> httpx.Response:

567 if input_type == "request":

568 path = "detection/v2/request-evaluations"

569 else:

570 path = "detection/v2/response-evaluations"

571

572 headers: Final = {

573 "Content-Type": "application/json",

574 "hl-runtime-edge-provider": "litellm",

575 "hl-runtime-edge-provider-version": "2",

576 }

577 if self.jwt_token:

578 headers["Authorization"] = f"Bearer {self.jwt_token}"

579

580 headers.update(hl_headers)

581

582 try:

583 response = await self._http_client.post(

584 f"{self.api_base}/{path}",

585 json=payload,

586 headers=headers,

587 )

588 response.raise_for_status()

589

590 verbose_proxy_logger.debug("Hiddenlayer reponse: %s", response)

591

592 return response

593 except HTTPStatusError as e:

594 # Try the request again by refreshing the jwt if we get 401

595 # since the Hiddenlayer jwt timeout is an hour and this is

596 # a long lived session application

597 if e.response.status_code == 401 and self.jwt_token is not None:

598 verbose_proxy_logger.debug(

599 "Unable to authenticate to Hiddenlayer, JWT token is invalid or expired, trying to refresh the token."

600 )

601 self.jwt_token = self.refresh_jwt_func()

602 headers["Authorization"] = f"Bearer {self.jwt_token}"

603 response = await self._http_client.post(

604 f"{self.api_base}/{path}",

605 json=payload,

606 headers=headers,

607 )

608 else:

609 raise e

610

611 response.raise_for_status()

612

613 verbose_proxy_logger.debug("Hiddenlayer reponse: %s", response)

614 return response

615

616 @staticmethod

617 def get_config_model() -> type[GuardrailConfigModel[BaseModel]] | None:

618 from litellm.types.proxy.guardrails.guardrail_hooks.hiddenlayer import (

619 HiddenlayerGuardrailConfigModel,

620 )

621

622 return HiddenlayerGuardrailConfigModel

Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py: 27%

278 statements