Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-guard/src/lexigram/ai/guard/di/provider.py: 67%

1"""GuardProtocol DI provider — registers guard pipeline with the container.""" 

2 

3from __future__ import annotations 

4 

5from typing import TYPE_CHECKING, Any, cast 

6 

7from lexigram.ai.guard.config import GuardConfig 

8from lexigram.ai.guard.input.injection import PromptInjectionDetector 

9from lexigram.ai.guard.input.length import InputLengthGuard 

10from lexigram.ai.guard.input.llm_injection import LLMInjectionDetector 

11from lexigram.ai.guard.input.llm_jailbreak import LLMJailbreakDetector 

12from lexigram.ai.guard.input.pii import PIIDetector 

13from lexigram.ai.guard.input.topic import TopicRestrictor 

14from lexigram.ai.guard.output.length import OutputLengthGuard 

15from lexigram.ai.guard.output.pii_redactor import PIIRedactor 

16from lexigram.ai.guard.pipeline.guard_pipeline import GuardPipeline 

17from lexigram.contracts.ai.guards import ( 

18 GuardPipelineProtocol, 

19 InputGuardProtocol, 

20 OutputGuardProtocol, 

21) 

22from lexigram.contracts.ai.llm import LLMClientProtocol 

23from lexigram.contracts.core.health import HealthCheckResult, HealthStatus 

24from lexigram.contracts.core.provider import ProviderPriority 

25from lexigram.di.provider import Provider 

26from lexigram.logging import ( 

27 get_logger, 

28) 

29 

30if TYPE_CHECKING: 

31 from lexigram.contracts.core.di import ( 

32 ContainerRegistrarProtocol, 

33 ContainerResolverProtocol, 

34 ) 

35 

36logger = get_logger(__name__) 

37 

38 

39class GuardProvider(Provider): 

40 """Provider for the content safety guard pipeline. 

41 

42 Reads :class:`~lexigram.ai.guard.config.GuardConfig`, builds a 

43 :class:`~lexigram.ai.guard.pipeline.guard_pipeline.GuardPipeline` 

44 with the configured guards, and registers it as a singleton. 

45 

46 When ``GuardConfig.enabled`` is ``False``, a no-op pipeline with no 

47 guards is registered (all checks pass through). 

48 

49 When ``GuardConfig.enable_llm_guards`` is ``True`` and an 

50 ``LLMClientProtocol`` is available in the container (resolved 

51 optionally during :meth:`boot`), LLM-based injection and jailbreak 

52 detectors are appended to the pipeline. 

53 """ 

54 

55 name = "guard" 

56 priority = ProviderPriority.SECURITY 

57 config_key: str | None = "ai_guard" 

58 config_model: type | None = GuardConfig 

59 

60 def __init__( 

61 self, 

62 config: GuardConfig | None = None, 

63 enable_audit_logging: bool = True, 

64 **kwargs: Any, 

65 ) -> None: 

66 super().__init__() 

67 self._requested_config = config 

68 self._config = config or GuardConfig() 

69 self._enable_audit_logging = enable_audit_logging 

70 self._pipeline: GuardPipeline | None = None 

71 

72 async def register(self, container: ContainerRegistrarProtocol) -> None: 

73 """Register the guard pipeline with the DI container.""" 

74 self._config = self._requested_config or ( 

75 self.config 

76 if isinstance(getattr(self, "config", None), GuardConfig) 

77 else self._config 

78 ) 

79 container.singleton(GuardConfig, self._config) 

80 

81 if not self._config.enabled: 

82 logger.info("guard_provider_disabled", reason="GuardConfig.enabled=False") 

83 pipeline = GuardPipeline() 

84 self._pipeline = pipeline 

85 container.singleton(GuardPipeline, pipeline) 

86 container.singleton(GuardPipelineProtocol, pipeline) 

87 return 

88 

89 pipeline = self._build_pipeline() 

90 self._pipeline = pipeline 

91 container.singleton(GuardPipeline, pipeline) 

92 container.singleton(GuardPipelineProtocol, pipeline) 

93 logger.info( 

94 "guard_provider_registered", 

95 input_guards=len(pipeline._input_guards), 

96 output_guards=len(pipeline._output_guards), 

97 ) 

98 

99 async def boot(self, container: ContainerResolverProtocol) -> None: 

100 """Boot phase — optionally attach LLM-based guards if configured. 

101 

102 Resolves ``LLMClientProtocol`` from the container when 

103 ``GuardConfig.enable_llm_guards`` is ``True``. If the client is 

104 not registered, LLM guards are silently skipped. 

105 """ 

106 if ( 

107 not self._config.enable_llm_guards 

108 or not self._config.enabled 

109 or self._pipeline is None 

110 ): 

111 logger.debug("guard_provider_booted", llm_guards=False) 

112 return 

113 

114 try: 

115 llm = await container.resolve_optional(LLMClientProtocol) 

116 except (ImportError, AttributeError, RuntimeError) as exc: 

117 logger.warning( 

118 "guard_provider_llm_resolve_failed", 

119 error=str(exc), 

120 ) 

121 logger.debug("guard_provider_booted", llm_guards=False) 

122 return 

123 

124 if llm is None: 

125 logger.info( 

126 "guard_provider_llm_guards_skipped", 

127 reason="LLMClientProtocol not registered in container", 

128 ) 

129 logger.debug("guard_provider_booted", llm_guards=False) 

130 return 

131 

132 cfg = self._config 

133 self._pipeline.add_input_guard( 

134 cast( 

135 "InputGuardProtocol", 

136 LLMInjectionDetector( 

137 llm, 

138 model=cfg.guard_model, 

139 threshold=cfg.llm_guard_threshold, 

140 action=cfg.injection_action, 

141 fail_open=cfg.llm_guard_fail_open, 

142 ), 

143 ) 

144 ) 

145 self._pipeline.add_input_guard( 

146 cast( 

147 "InputGuardProtocol", 

148 LLMJailbreakDetector( 

149 llm, 

150 model=cfg.guard_model, 

151 threshold=cfg.llm_guard_threshold, 

152 action=cfg.injection_action, 

153 fail_open=cfg.llm_guard_fail_open, 

154 ), 

155 ) 

156 ) 

157 logger.info( 

158 "guard_provider_llm_guards_registered", 

159 model=cfg.guard_model, 

160 threshold=cfg.llm_guard_threshold, 

161 ) 

162 logger.debug("guard_provider_booted", llm_guards=True) 

163 

164 async def shutdown(self) -> None: 

165 """Shutdown phase — no cleanup required for guard pipeline.""" 

166 logger.debug("guard_provider_shutdown") 

167 

168 async def health_check(self, timeout: float = 5.0) -> HealthCheckResult: 

169 """Health check — always healthy (in-process domain provider). 

170 

171 No external backend to ping. 

172 

173 Args: 

174 timeout: Ignored for in-process providers. 

175 

176 Returns: 

177 Always HEALTHY — no external backend to ping. 

178 """ 

179 return HealthCheckResult( 

180 component=self.name, 

181 status=HealthStatus.HEALTHY, 

182 details={"status": "operational"}, 

183 ) 

184 

185 def _build_pipeline(self) -> GuardPipeline: 

186 """Build the heuristic guard pipeline from configuration. 

187 

188 Returns: 

189 Configured GuardPipeline instance. 

190 """ 

191 cfg = self._config 

192 

193 # ---------- input guards ---------- 

194 input_guards: list[InputGuardProtocol] = [] 

195 

196 if cfg.injection_detection: 

197 input_guards.append(PromptInjectionDetector(action=cfg.injection_action)) 

198 

199 if cfg.pii_detection: 

200 entities = cfg.pii_entities or None 

201 input_guards.append(PIIDetector(action=cfg.pii_action, entities=entities)) 

202 

203 if cfg.max_input_chars > 0: 

204 input_guards.append( 

205 InputLengthGuard( 

206 max_chars=cfg.max_input_chars, action=cfg.length_action 

207 ) 

208 ) 

209 

210 if cfg.restricted_topics: 

211 input_guards.append( 

212 TopicRestrictor(restricted_topics=cfg.restricted_topics) 

213 ) 

214 

215 # ---------- output guards ---------- 

216 output_guards: list[OutputGuardProtocol] = [] 

217 

218 if cfg.pii_redaction_output: 

219 entities = cfg.pii_entities or None 

220 output_guards.append(PIIRedactor(entities=entities)) 

221 

222 if cfg.max_output_chars > 0: 

223 output_guards.append( 

224 OutputLengthGuard( 

225 max_chars=cfg.max_output_chars, action=cfg.length_action 

226 ) 

227 ) 

228 

229 return GuardPipeline(input_guards=input_guards, output_guards=output_guards) 

230 

231 

232__all__ = ["GuardProvider"]