Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py: 22%

1""" 

2Custom code guardrail for LiteLLM. 

3 

4This module provides a guardrail that executes user-defined Python-like code 

5to implement custom guardrail logic. The code runs in a sandboxed environment 

6with access to LiteLLM-provided primitives for common guardrail operations. 

7 

8Example custom code (sync): 

9 

10 def apply_guardrail(inputs, request_data, input_type): 

11 '''Block messages containing SSNs''' 

12 for text in inputs["texts"]: 

13 if regex_match(text, r"\\d{3}-\\d{2}-\\d{4}"): 

14 return block("Social Security Number detected") 

15 return allow() 

16 

17Example custom code (async with HTTP): 

18 

19 async def apply_guardrail(inputs, request_data, input_type): 

20 '''Call external moderation API''' 

21 for text in inputs["texts"]: 

22 response = await http_post( 

23 "https://api.example.com/moderate", 

24 body={"text": text} 

25 ) 

26 if response["success"] and response["body"].get("flagged"): 

27 return block("Content flagged by moderation API") 

28 return allow() 

29 

30Example: block when response rejects the user (input_type response only): 

31 

32 Use RESPONSE_REJECTION_GUARDRAIL_CODE from .response_rejection_code — it 

33 checks response texts for phrases like "That's not something I can help with" 

34 and returns block() so the guardrail raises a block error. 

35""" 

36 

37import asyncio 

38import threading 

39import time 

40from collections.abc import Callable, Mapping 

41from typing import TYPE_CHECKING, Final, Literal, Optional, cast 

42 

43from fastapi import HTTPException 

44from typing_extensions import TypedDict, Unpack 

45 

46from litellm._logging import verbose_proxy_logger 

47from litellm.exceptions import ModifyResponseException 

48from litellm.integrations.custom_guardrail import ( 

49 CustomGuardrail, 

50 log_guardrail_information, 

51) 

52from litellm.types.guardrails import GuardrailEventHooks 

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

54from litellm.types.utils import GenericGuardrailAPIInputs 

55 

56from .sandbox import build_sandbox_globals, compile_sandboxed 

57 

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

59 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj 

60 

61 

62def _metadata_bucket(request_data: Mapping[str, object], key: str) -> Mapping[str, object]: 

63 bucket: Final = request_data.get(key) 

64 return bucket if isinstance(bucket, Mapping) else {} 

65 

66 

67class CustomCodeGuardrailError(Exception): 

68 """Raised when custom code guardrail execution fails.""" 

69 

70 def __init__(self, message: str, details: Mapping[str, object] | None = None) -> None: 

71 super().__init__(message) 

72 self.details: Mapping[str, object] = details or {} 

73 

74 

75class CustomCodeCompilationError(CustomCodeGuardrailError): 

76 """Raised when custom code fails to compile.""" 

77 

78 

79class CustomCodeExecutionError(CustomCodeGuardrailError): 

80 """Raised when custom code fails during execution.""" 

81 

82 

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

84 """Base-class constructor options this guardrail forwards untouched to CustomGuardrail.""" 

85 

86 

87class CustomCodeGuardrailConfigModel(GuardrailConfigModel): 

88 """Configuration parameters for the custom code guardrail.""" 

89 

90 custom_code: str 

91 """The Python-like code containing the apply_guardrail function.""" 

92 

93 

94class CustomCodeGuardrail(CustomGuardrail): 

95 """ 

96 Guardrail that executes user-defined Python-like code. 

97 

98 The code runs in a sandboxed environment that provides: 

99 - Access to LiteLLM primitives (regex_match, json_parse, etc.) 

100 - No file I/O or network access 

101 - No imports allowed 

102 

103 Users write an `apply_guardrail(inputs, request_data, input_type)` function 

104 that returns one of: 

105 - allow() - let the request/response through 

106 - block(reason) - reject with a message 

107 - flag(reason) - let it through but log a non-blocking violation 

108 - modify(texts=...) - transform the content 

109 

110 Example: 

111 def apply_guardrail(inputs, request_data, input_type): 

112 for text in inputs["texts"]: 

113 if regex_match(text, r"password"): 

114 return block("Sensitive content detected") 

115 return allow() 

116 """ 

117 

118 def __init__( 

119 self, 

120 custom_code: str, 

121 guardrail_name: str | None = "custom_code", 

122 **kwargs: Unpack[_CustomGuardrailOptions], 

123 ) -> None: 

124 """ 

125 Initialize the custom code guardrail. 

126 

127 Args: 

128 custom_code: The source code containing apply_guardrail function 

129 guardrail_name: Name of this guardrail instance 

130 **kwargs: Additional arguments passed to CustomGuardrail 

131 """ 

132 self.custom_code: str = custom_code 

133 self._compiled_function: Callable[..., object] | None = None 

134 self._compile_lock = threading.Lock() 

135 self._compile_error: str | None = None 

136 

137 super().__init__( 

138 guardrail_name=guardrail_name, 

139 supported_event_hooks=list(self.get_supported_event_hooks()), 

140 **kwargs, 

141 ) 

142 

143 # Compile the code on initialization 

144 self._compile_custom_code() 

145 

146 @staticmethod 

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

148 """Returns the config model for the UI.""" 

149 return CustomCodeGuardrailConfigModel 

150 

151 @classmethod 

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

153 return [ 

154 GuardrailEventHooks.pre_call, 

155 GuardrailEventHooks.during_call, 

156 GuardrailEventHooks.post_call, 

157 GuardrailEventHooks.pre_mcp_call, 

158 GuardrailEventHooks.during_mcp_call, 

159 GuardrailEventHooks.logging_only, 

160 ] 

161 

162 def _do_compile(self) -> None: 

163 """Internal compilation method without lock. Expected to run inside _compile_lock.""" 

164 exec_globals: Final = build_sandbox_globals() 

165 compiled: Final = compile_sandboxed(self.custom_code) 

166 exec(compiled, exec_globals) # noqa: S102 

167 

168 if "apply_guardrail" not in exec_globals: 

169 raise CustomCodeCompilationError( 

170 "Custom code must define an 'apply_guardrail' function. " 

171 "Expected signature: apply_guardrail(inputs, request_data, input_type)" 

172 ) 

173 

174 apply_fn: Final = exec_globals["apply_guardrail"] 

175 if not callable(apply_fn): 

176 raise CustomCodeCompilationError("'apply_guardrail' must be a callable function") 

177 

178 self._compiled_function = apply_fn 

179 

180 def _compile_custom_code(self) -> None: 

181 """ 

182 Compile the custom code and extract the apply_guardrail function. 

183 

184 The code runs in a sandboxed environment with only the allowed primitives. 

185 """ 

186 with self._compile_lock: 

187 if self._compiled_function is not None: 

188 return 

189 

190 try: 

191 self._do_compile() 

192 verbose_proxy_logger.debug("Custom code guardrail '%s' compiled successfully", self.guardrail_name) 

193 

194 except SyntaxError as e: 

195 self._compile_error = f"Syntax error in custom code: {e}" 

196 raise CustomCodeCompilationError(self._compile_error) from e 

197 except CustomCodeCompilationError: 

198 raise 

199 except Exception as e: 

200 self._compile_error = f"Failed to compile custom code: {e}" 

201 raise CustomCodeCompilationError(self._compile_error) from e 

202 

203 @log_guardrail_information 

204 async def apply_guardrail( 

205 self, 

206 inputs: GenericGuardrailAPIInputs, 

207 request_data: dict[str, object], 

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

209 logging_obj: Optional["LiteLLMLoggingObj"] = None, 

210 ) -> GenericGuardrailAPIInputs: 

211 """ 

212 Apply the custom code guardrail to the inputs. 

213 

214 This method calls the user-defined apply_guardrail function and 

215 processes its result to determine the appropriate action. 

216 

217 The user-defined function can be either sync or async: 

218 - Sync: def apply_guardrail(inputs, request_data, input_type): ... 

219 - Async: async def apply_guardrail(inputs, request_data, input_type): ... 

220 

221 Async functions are recommended when using http_request, http_get, or 

222 http_post primitives to avoid blocking the event loop. 

223 

224 Args: 

225 inputs: Dictionary containing texts, images, tool_calls 

226 request_data: The original request data with metadata 

227 input_type: "request" for pre-call, "response" for post-call 

228 logging_obj: Optional logging object 

229 

230 Returns: 

231 GenericGuardrailAPIInputs - possibly modified 

232 

233 Raises: 

234 HTTPException: If content is blocked 

235 CustomCodeExecutionError: If execution fails 

236 """ 

237 if self._compiled_function is None: 

238 if self._compile_error: 

239 raise CustomCodeExecutionError(f"Custom code guardrail not compiled: {self._compile_error}") 

240 raise CustomCodeExecutionError("Custom code guardrail not compiled") 

241 

242 start_time: Final = time.time() 

243 try: 

244 # Prepare inputs dict for the function 

245 

246 # Prepare request_data with safe subset of information 

247 safe_request_data: Final = self._prepare_safe_request_data(request_data) 

248 

249 # Execute the custom function - handle both sync and async functions 

250 raw_result: Final = self._compiled_function(inputs, safe_request_data, input_type) 

251 

252 # If the function is async (returns a coroutine), await it 

253 resolved_result: Final[object] = await raw_result if asyncio.iscoroutine(raw_result) else raw_result 

254 

255 # Process the result 

256 return self._process_result( 

257 result=resolved_result, 

258 inputs=inputs, 

259 request_data=request_data, 

260 input_type=input_type, 

261 start_time=start_time, 

262 ) 

263 

264 except HTTPException: 

265 # Re-raise HTTP exceptions (from block action) 

266 raise 

267 except ModifyResponseException: 

268 # Pre-call block uses passthrough; must not wrap as execution error (500) 

269 raise 

270 except Exception as e: 

271 verbose_proxy_logger.error("Custom code guardrail '%s' execution error: %s", self.guardrail_name, e) 

272 raise CustomCodeExecutionError( 

273 f"Custom code guardrail execution failed: {e}", 

274 details={ 

275 "guardrail_name": self.guardrail_name, 

276 "input_type": input_type, 

277 }, 

278 ) from e 

279 

280 def _prepare_safe_request_data(self, request_data: Mapping[str, object]) -> dict[str, object]: 

281 """ 

282 Prepare a safe subset of request_data for code execution. 

283 

284 This filters out sensitive information and provides only what's 

285 needed for guardrail logic. 

286 

287 Args: 

288 request_data: The full request data 

289 

290 Returns: 

291 Safe subset of request data 

292 """ 

293 metadata: Final = { 

294 **_metadata_bucket(request_data, "metadata"), 

295 **_metadata_bucket(request_data, "litellm_metadata"), 

296 } 

297 return { 

298 "model": request_data.get("model"), 

299 "user_id": metadata.get("user_api_key_user_id"), 

300 "team_id": metadata.get("user_api_key_team_id"), 

301 "end_user_id": metadata.get("user_api_key_end_user_id"), 

302 "metadata": metadata, 

303 } 

304 

305 def _process_result( 

306 self, 

307 result: object, 

308 inputs: GenericGuardrailAPIInputs, 

309 request_data: dict[str, object], 

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

311 start_time: float, 

312 ) -> GenericGuardrailAPIInputs: 

313 """ 

314 Process the result from the custom code function. 

315 

316 Args: 

317 result: The return value from apply_guardrail 

318 inputs: The original inputs 

319 request_data: The request data 

320 input_type: "request" or "response" 

321 start_time: Unix timestamp of when the guardrail started running, used for the flagged log entry 

322 

323 Returns: 

324 GenericGuardrailAPIInputs - possibly modified 

325 

326 Raises: 

327 HTTPException: If action is "block" 

328 """ 

329 if not isinstance(result, dict): 

330 verbose_proxy_logger.warning( 

331 "Custom code guardrail '%s': Expected dict result, got %s. Treating as allow.", 

332 self.guardrail_name, 

333 type(result).__name__, 

334 ) 

335 return inputs 

336 

337 action: Final = result.get("action", "allow") 

338 

339 if action == "allow": 

340 verbose_proxy_logger.debug("Custom code guardrail '%s': Allowing %s", self.guardrail_name, input_type) 

341 return inputs 

342 

343 elif action == "block": 

344 reason: Final = result.get("reason", "Blocked by custom code guardrail") 

345 detection_info: Final = result.get("detection_info", {}) 

346 

347 verbose_proxy_logger.info( 

348 "Custom code guardrail '%s': Blocking %s - %s", self.guardrail_name, input_type, reason 

349 ) 

350 

351 is_output: Final = input_type == "response" 

352 

353 # For pre-call, raise passthrough exception to return synthetic response 

354 if not is_output: 

355 self.raise_passthrough_exception( 

356 violation_message=reason, 

357 request_data=request_data, 

358 detection_info=detection_info, 

359 ) 

360 

361 # For post-call, raise HTTP exception 

362 raise HTTPException( 

363 status_code=400, 

364 detail={ 

365 "error": reason, 

366 "guardrail": self.guardrail_name, 

367 "detection_info": detection_info, 

368 }, 

369 ) 

370 

371 elif action == "flag": 

372 flag_reason: Final = result.get("reason", "Flagged by custom code guardrail") 

373 verbose_proxy_logger.info( 

374 "Custom code guardrail '%s': Flagging %s - %s", self.guardrail_name, input_type, flag_reason 

375 ) 

376 end_time: Final = time.time() 

377 self.add_standard_logging_guardrail_information_to_request_data( 

378 guardrail_json_response={ # mutable-ok: logging helper requires a dict 

379 "action": "flag", 

380 "reason": flag_reason, 

381 "input_type": input_type, 

382 "metadata": result.get("metadata") or {}, # mutable-ok: logging helper requires a dict 

383 }, 

384 request_data=request_data, 

385 guardrail_status="guardrail_flagged", 

386 start_time=start_time, 

387 end_time=end_time, 

388 duration=end_time - start_time, 

389 ) 

390 return inputs 

391 

392 elif action == "modify": 

393 verbose_proxy_logger.debug("Custom code guardrail '%s': Modifying %s", self.guardrail_name, input_type) 

394 

395 # Apply modifications 

396 modified_inputs: Final = dict(inputs) 

397 

398 if "texts" in result and result["texts"] is not None: 

399 modified_inputs["texts"] = result["texts"] 

400 

401 if "images" in result and result["images"] is not None: 

402 modified_inputs["images"] = result["images"] 

403 

404 if "tool_calls" in result and result["tool_calls"] is not None: 

405 modified_inputs["tool_calls"] = result["tool_calls"] 

406 

407 return cast(GenericGuardrailAPIInputs, modified_inputs) 

408 

409 else: 

410 verbose_proxy_logger.warning( 

411 "Custom code guardrail '%s': Unknown action '%s'. Treating as allow.", self.guardrail_name, action 

412 ) 

413 return inputs 

414 

415 def update_custom_code(self, new_code: str) -> None: 

416 """ 

417 Update the custom code and recompile. 

418 

419 This method allows hot-reloading of guardrail logic without 

420 restarting the server. 

421 

422 Args: 

423 new_code: The new source code 

424 

425 Raises: 

426 CustomCodeCompilationError: If the new code fails to compile 

427 """ 

428 with self._compile_lock: 

429 # Reset state 

430 old_function: Final = self._compiled_function 

431 old_code: Final = self.custom_code 

432 self._compiled_function = None 

433 self._compile_error = None 

434 

435 try: 

436 self.custom_code = new_code 

437 self._do_compile() 

438 verbose_proxy_logger.info("Custom code guardrail '%s': Code updated successfully", self.guardrail_name) 

439 except SyntaxError as e: 

440 # Rollback on failure 

441 self.custom_code = old_code 

442 self._compiled_function = old_function 

443 self._compile_error = f"Syntax error in custom code: {e}" 

444 raise CustomCodeCompilationError(self._compile_error) from e 

445 except CustomCodeCompilationError: 

446 # Rollback on failure 

447 self.custom_code = old_code 

448 self._compiled_function = old_function 

449 raise 

450 except Exception as e: 

451 # Rollback on failure 

452 self.custom_code = old_code 

453 self._compiled_function = old_function 

454 self._compile_error = f"Failed to compile custom code: {e}" 

455 raise CustomCodeCompilationError(self._compile_error) from e