Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-guard/src/lexigram/ai/guard/pipeline/guard_pipeline.py: 22%

87 statements  

« prev     ^ index     » next       coverage.py v7.15.4, created at 2026-08-25 07:19 +0800

1"""GuardProtocol pipeline — orchestrates input and output guards. 

2 

3The pipeline runs guards in registration order. For input guards it 

4applies them sequentially to the request content; for output guards it 

5applies them to the LLM response. Any BLOCK result stops further 

6evaluation and returns immediately. 

7""" 

8 

9from __future__ import annotations 

10 

11from typing import TYPE_CHECKING, Any, TypeAlias 

12 

13from lexigram.ai.guard.pipeline.result import ( 

14 AggregateGuardResult, 

15 GuardAction, 

16) 

17from lexigram.contracts.ai.exceptions import GuardError 

18from lexigram.contracts.ai.guards import GuardResultProtocol 

19from lexigram.logging import ( 

20 get_logger, 

21) 

22from lexigram.result import Err, Ok, Result 

23 

24if TYPE_CHECKING: 

25 from lexigram.contracts.ai.guards import InputGuardProtocol, OutputGuardProtocol 

26 

27logger = get_logger(__name__) 

28 

29GuardResult: TypeAlias = Result[GuardResultProtocol, GuardError] 

30GuardResultOrException: TypeAlias = GuardResult | BaseException 

31 

32 

33class GuardPipeline: 

34 """Ordered chain of input and output guards. 

35 

36 Guards are evaluated in order. For input, content flows through 

37 each guard; if a guard redacts content, the redacted version is 

38 passed to subsequent guards. A BLOCK stops evaluation immediately. 

39 

40 Args: 

41 input_guards: Guards to run on user inputs before sending to LLM. 

42 output_guards: Guards to run on LLM outputs before returning to caller. 

43 

44 Example:: 

45 

46 pipeline = GuardPipeline( 

47 input_guards=[ 

48 PromptInjectionDetector(action="block"), 

49 PIIDetector(action="redact"), 

50 ], 

51 output_guards=[ 

52 PIIRedactor(entities=["SSN", "CREDIT_CARD"]), 

53 LengthGuard(max_chars=10000, action="block"), 

54 ], 

55 ) 

56 

57 result = await pipeline.check_input(user_message) 

58 if result.blocked: 

59 return Err(GuardViolationError(result)) 

60 safe_input = result.final_content or user_message 

61 """ 

62 

63 def __init__( 

64 self, 

65 input_guards: list[InputGuardProtocol] | None = None, 

66 output_guards: list[OutputGuardProtocol] | None = None, 

67 ) -> None: 

68 """Initialise the guard pipeline. 

69 

70 Args: 

71 input_guards: Ordered list of input guards to apply. 

72 output_guards: Ordered list of output guards to apply. 

73 """ 

74 self._input_guards: list[InputGuardProtocol] = input_guards or [] 

75 self._output_guards: list[OutputGuardProtocol] = output_guards or [] 

76 

77 def add_input_guard(self, guard: InputGuardProtocol) -> None: 

78 """Append an input guard to the pipeline. 

79 

80 Args: 

81 guard: GuardProtocol to add. 

82 """ 

83 self._input_guards.append(guard) 

84 

85 def add_output_guard(self, guard: OutputGuardProtocol) -> None: 

86 """Append an output guard to the pipeline. 

87 

88 Args: 

89 guard: GuardProtocol to add. 

90 """ 

91 self._output_guards.append(guard) 

92 

93 async def check_input( 

94 self, 

95 content: str, 

96 *, 

97 messages: list[Any] | None = None, 

98 metadata: dict[str, Any] | None = None, 

99 parallel: bool = False, 

100 ) -> Result[AggregateGuardResult, GuardError]: 

101 """Run all input guards against the content. 

102 

103 Guards are applied in order. Redacted content is forwarded to 

104 subsequent guards. A BLOCK terminates evaluation immediately. 

105 If `parallel` is True, all guards are run concurrently and redaction 

106 chaining is disabled (each operates on the original content). 

107 

108 Args: 

109 content: Raw input text to evaluate. 

110 messages: Optional structured message list for context. 

111 metadata: Optional request metadata (user_id, model, etc.). 

112 parallel: Whether to run guards concurrently with asyncio.gather. 

113 

114 Returns: 

115 Ok(AggregateGuardResult) combining all individual guard outcomes, 

116 or Err(GuardError) if a guard fails unexpectedly. 

117 """ 

118 if not self._input_guards: 

119 return Ok( 

120 AggregateGuardResult( 

121 passed=True, 

122 action=GuardAction.PASS, 

123 results=[], 

124 final_content=content, 

125 ) 

126 ) 

127 

128 results: list[GuardResultProtocol] = [] 

129 current_content = content 

130 

131 if parallel: 

132 import asyncio 

133 

134 coros = [ 

135 guard.check(content, messages=messages, metadata=metadata) 

136 for guard in self._input_guards 

137 ] 

138 check_results: list[GuardResultOrException] = await asyncio.gather( 

139 *coros, return_exceptions=True 

140 ) 

141 

142 for _guard, check_result in zip( 

143 self._input_guards, check_results, strict=False 

144 ): 

145 if isinstance(check_result, BaseException): 

146 return Err( 

147 GuardError(f"Guard failed: {check_result}", cause=check_result) 

148 ) 

149 

150 if check_result.is_err(): 

151 return Err(check_result.unwrap_err()) 

152 

153 res = check_result.unwrap() 

154 results.append(res) 

155 logger.debug( 

156 "input_guard_evaluated_parallel", 

157 guard=res.guard_name, 

158 action=res.action, 

159 ) 

160 

161 if res.action == GuardAction.BLOCK: 

162 logger.info( 

163 "input_guard_blocked", 

164 guard=res.guard_name, 

165 reason=res.details.get("reason", ""), 

166 ) 

167 return Ok( 

168 AggregateGuardResult( 

169 passed=False, 

170 action=GuardAction.BLOCK, 

171 results=results, 

172 final_content=None, 

173 ) 

174 ) 

175 

176 return Ok(AggregateGuardResult.from_results(results, content)) 

177 

178 for guard in self._input_guards: 

179 guard_result = await guard.check( 

180 current_content, 

181 messages=messages, 

182 metadata=metadata, 

183 ) 

184 

185 if guard_result.is_err(): 

186 return Err(guard_result.unwrap_err()) 

187 

188 res = guard_result.unwrap() 

189 results.append(res) 

190 logger.debug( 

191 "input_guard_evaluated", 

192 guard=res.guard_name, 

193 action=res.action, 

194 ) 

195 

196 if res.action == GuardAction.BLOCK: 

197 logger.info( 

198 "input_guard_blocked", 

199 guard=res.guard_name, 

200 reason=res.details.get("reason", ""), 

201 ) 

202 return Ok( 

203 AggregateGuardResult( 

204 passed=False, 

205 action=GuardAction.BLOCK, 

206 results=results, 

207 final_content=None, 

208 ) 

209 ) 

210 

211 if res.action == GuardAction.REDACT and res.redacted_content is not None: 

212 current_content = res.redacted_content 

213 

214 return Ok(AggregateGuardResult.from_results(results, current_content)) 

215 

216 async def check_output( 

217 self, 

218 content: str, 

219 *, 

220 original_input: str | None = None, 

221 metadata: dict[str, Any] | None = None, 

222 parallel: bool = False, 

223 ) -> Result[AggregateGuardResult, GuardError]: 

224 """Run all output guards against the LLM response. 

225 

226 Args: 

227 content: LLM response text to evaluate. 

228 original_input: The original user input for context. 

229 metadata: Optional request metadata (model, provider, etc.). 

230 parallel: Whether to run guards concurrently. 

231 

232 Returns: 

233 Ok(AggregateGuardResult) combining all individual guard outcomes. 

234 """ 

235 if not self._output_guards: 

236 return Ok( 

237 AggregateGuardResult( 

238 passed=True, 

239 action=GuardAction.PASS, 

240 results=[], 

241 final_content=content, 

242 ) 

243 ) 

244 

245 results: list[GuardResultProtocol] = [] 

246 current_content = content 

247 

248 if parallel: 

249 import asyncio 

250 

251 coros = [ 

252 guard.check(content, original_input=original_input, metadata=metadata) 

253 for guard in self._output_guards 

254 ] 

255 check_results: list[GuardResultOrException] = await asyncio.gather( 

256 *coros, return_exceptions=True 

257 ) 

258 

259 for _guard, check_result in zip( 

260 self._output_guards, check_results, strict=False 

261 ): 

262 if isinstance(check_result, BaseException): 

263 return Err( 

264 GuardError(f"Guard failed: {check_result}", cause=check_result) 

265 ) 

266 

267 if check_result.is_err(): 

268 return Err(check_result.unwrap_err()) 

269 

270 res = check_result.unwrap() 

271 results.append(res) 

272 logger.debug( 

273 "output_guard_evaluated_parallel", 

274 guard=res.guard_name, 

275 action=res.action, 

276 ) 

277 

278 if res.action == GuardAction.BLOCK: 

279 logger.info( 

280 "output_guard_blocked", 

281 guard=res.guard_name, 

282 reason=res.details.get("reason", ""), 

283 ) 

284 return Ok( 

285 AggregateGuardResult( 

286 passed=False, 

287 action=GuardAction.BLOCK, 

288 results=results, 

289 final_content=None, 

290 ) 

291 ) 

292 

293 return Ok(AggregateGuardResult.from_results(results, content)) 

294 

295 for guard in self._output_guards: 

296 guard_result = await guard.check( 

297 current_content, 

298 original_input=original_input, 

299 metadata=metadata, 

300 ) 

301 if guard_result.is_err(): 

302 return Err(guard_result.unwrap_err()) 

303 

304 res = guard_result.unwrap() 

305 results.append(res) 

306 logger.debug( 

307 "output_guard_evaluated", 

308 guard=res.guard_name, 

309 action=res.action, 

310 ) 

311 

312 if res.action == GuardAction.BLOCK: 

313 logger.info( 

314 "output_guard_blocked", 

315 guard=res.guard_name, 

316 reason=res.details.get("reason", ""), 

317 ) 

318 return Ok( 

319 AggregateGuardResult( 

320 passed=False, 

321 action=GuardAction.BLOCK, 

322 results=results, 

323 final_content=None, 

324 ) 

325 ) 

326 

327 if res.action == GuardAction.REDACT and res.redacted_content is not None: 

328 current_content = res.redacted_content 

329 

330 return Ok(AggregateGuardResult.from_results(results, current_content)) 

331 

332 

333__all__ = ["GuardPipeline"]