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"]