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

1"""LLM Provider for Lexigram Framework dependency injection. 

2 

3Registers all LLM-layer services (clients, cache, model manager) with 

4the container so they can be injected throughout the application. 

5""" 

6 

7from __future__ import annotations 

8 

9from typing import TYPE_CHECKING, cast 

10 

11from lexigram.ai.llm.exceptions import LLMError 

12from lexigram.contracts.ai import LLMClientProtocol 

13from lexigram.contracts.ai.providers import ProviderRegistryProtocol 

14from lexigram.contracts.core.health import HealthCheckResult 

15from lexigram.contracts.exceptions.base import LexigramError 

16from lexigram.contracts.exceptions.container import UnresolvableDependencyError 

17from lexigram.contracts.exceptions.provider import ModuleVisibilityError 

18from lexigram.di.decorators import inject 

19from lexigram.di.provider import Provider, ProviderPriority 

20from lexigram.logging import get_logger 

21 

22if TYPE_CHECKING: 

23 from lexigram.contracts.core.di import ( 

24 ContainerRegistrarProtocol, 

25 ContainerResolverProtocol, 

26 ) 

27 from lexigram.contracts.core.health import HealthCheckResult 

28 from lexigram.contracts.infra.cache import CacheBackendProtocol 

29 

30from lexigram.ai.llm.config import ClientConfig 

31from lexigram.ai.llm.di.factories import create_llm_cache, create_llm_client 

32from lexigram.ai.llm.model_manager import LLMModelManager 

33from lexigram.ai.llm.registry.core import ProviderRegistry 

34 

35logger = get_logger(__name__) 

36 

37__all__ = ["LLMProvider"] 

38 

39 

40# Map provider name → pyproject.toml extra name for actionable install hints. 

41# Most match the provider name, but some share an extra (e.g. deepseek, azure-openai 

42# all bundle under the ``openai`` extra). 

43_PROVIDER_EXTRAS: dict[str, str] = { 

44 "ollama": "ollama", 

45 "openai": "openai", 

46 "anthropic": "anthropic", 

47 "cohere": "cohere", 

48 "groq": "groq", 

49 "mistral": "mistral", 

50 "deepseek": "openai", 

51 "together": "openai", 

52 "fireworks": "openai", 

53 "openrouter": "openai", 

54 "azure-openai": "openai", 

55 "cloudflare": "openai", 

56 "gemini": "openai", 

57 "google-vertex": "openai", 

58 "aws-bedrock": "openai", 

59} 

60 

61 

62@inject 

63class LLMProvider(Provider): 

64 """Provider that registers LLM services with the Lexigram DI container. 

65 

66 Registers an LLMClientProtocol, optional LLM response cache, and an LLMModelManager 

67 so all three are injectable throughout the application. 

68 

69 Example: 

70 >>> from lexigram.ai.llm.di.provider import LLMProvider 

71 >>> from lexigram.ai.llm.config import ClientConfig 

72 >>> 

73 >>> app.use(LLMProvider(ClientConfig(provider="openai", model="gpt-4o"))) 

74 >>> 

75 >>> # LLMClientProtocol is now injectable: 

76 >>> class MyService: 

77 ... def __init__(self, llm: LLMClientProtocol) -> None: 

78 ... self.llm = llm 

79 """ 

80 

81 name = "llm" 

82 priority = ProviderPriority.DOMAIN 

83 config_key: str | None = "ai_llm" 

84 config_model: type | None = ClientConfig 

85 

86 def __init__( 

87 self, 

88 config: ClientConfig | None = None, 

89 enable_model_manager: bool = False, 

90 enable_streaming: bool = True, 

91 audit_calls: bool = False, 

92 name: str = "llm", 

93 cache_backend: CacheBackendProtocol | None = None, 

94 stub_mode: bool = False, 

95 ) -> None: 

96 """Initialize the LLM Provider. 

97 

98 Args: 

99 config: LLM client configuration; defaults to ClientConfig() (reads env). 

100 enable_model_manager: Register LLMModelManager for local model control. 

101 enable_streaming: Enable streaming response support. 

102 audit_calls: Register :class:`~lexigram.ai.llm.audit_bridge.LLMAuditBridge` 

103 to emit audit entries per completion. 

104 name: Provider name used for identification. 

105 cache_backend: Injected cache backend for optional response caching. 

106 """ 

107 super().__init__(name=name) 

108 self._requested_config = config 

109 self.config = config or ClientConfig() 

110 self.enable_model_manager = enable_model_manager 

111 self.enable_streaming = enable_streaming 

112 self.audit_calls = audit_calls 

113 self.cache_backend = cache_backend 

114 self.stub_mode = stub_mode 

115 self._llm_client: LLMClientProtocol | None = None 

116 

117 @staticmethod 

118 def _instructor_available() -> bool: 

119 """Check if the instructor library is available.""" 

120 import importlib.util 

121 

122 try: 

123 return importlib.util.find_spec("instructor") is not None 

124 except (ValueError, AttributeError): 

125 return False 

126 

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

128 """Register LLM services with the DI container. 

129 

130 Args: 

131 container: The Lexigram DI container registrar. 

132 """ 

133 from lexigram.ai.llm.pricing.registry import TokenCounterRegistry 

134 from lexigram.contracts.ai.llm import TokenCounterProtocol 

135 

136 self.config = self._requested_config or self.config 

137 container.singleton(ClientConfig, self.config) 

138 

139 if not self.config.enabled: 

140 logger.info("llm_disabled", reason="ClientConfig.enabled=False") 

141 return 

142 

143 logger.info( 

144 "Registering LLM services", 

145 provider=self.config.provider, 

146 model=self.config.model, 

147 ) 

148 

149 registry = ProviderRegistry() 

150 container.singleton(ProviderRegistry, registry) 

151 container.singleton(ProviderRegistryProtocol, registry) 

152 

153 # Register TokenCounterRegistry 

154 token_registry = TokenCounterRegistry.with_defaults() 

155 container.singleton(TokenCounterRegistry, token_registry) 

156 

157 # Register ParserRegistry with default parsers 

158 from lexigram.ai.llm.parsers.registry import ParserRegistry 

159 

160 parser_registry = ParserRegistry.with_defaults() 

161 container.singleton(ParserRegistry, parser_registry) 

162 logger.info("Registered ParserRegistry with default parsers") 

163 

164 llm_client: LLMClientProtocol 

165 try: 

166 if self.stub_mode: 

167 from lexigram.ai.llm.clients.noop import NoOpLLMClient 

168 

169 llm_client = cast("LLMClientProtocol", NoOpLLMClient(self.config)) 

170 else: 

171 llm_client = await create_llm_client(self.config, registry) 

172 except ImportError as exc: 

173 extra = _PROVIDER_EXTRAS.get( 

174 self.config.provider.value, 

175 self.config.provider.value, 

176 ) 

177 msg = ( 

178 f"Missing SDK for provider {self.config.provider!r}. " 

179 f"The SDK package is an optional dependency.\n" 

180 f" Install: uv sync --extra {extra}\n" 

181 f" Or: pip install lexigram-ai-llm[{extra}]\n" 

182 f" Error: {exc}" 

183 ) 

184 raise LLMError(msg) from exc 

185 

186 # Always enrich completions with provenance (provider, model_revision, 

187 # prompt_hash) — innermost wrap so downstream layers see fully-populated 

188 # Completion objects. LXF-003 wiring. 

189 from lexigram.ai.llm.wrappers import CompletionEnricher, LLMCacheWrapper 

190 

191 llm_client = CompletionEnricher.wrap( 

192 llm_client, 

193 provider=self.config.provider, 

194 model=self.config.model, 

195 model_revision=getattr(self.config, "model_revision", None), 

196 ) 

197 

198 container.singleton(LLMClientProtocol, llm_client) 

199 container.singleton("llm", llm_client) 

200 self._llm_client = llm_client 

201 logger.info( 

202 "Registered LLM client", 

203 provider=self.config.provider, 

204 model=self.config.model, 

205 ) 

206 

207 if self.config.enable_cache: 

208 llm_cache = await create_llm_cache(self.config, self.cache_backend) 

209 container.singleton("llm_cache", llm_cache) 

210 logger.info("Registered LLM response cache") 

211 

212 if self.enable_model_manager: 

213 container.singleton(LLMModelManager) 

214 container.singleton("llm_model_manager", LLMModelManager) 

215 logger.info("Registered LLMModelManager") 

216 

217 if self.audit_calls: 

218 from lexigram.ai.llm.audit_bridge import LLMAuditBridge 

219 

220 container.singleton(LLMAuditBridge) 

221 wrapped_client = LLMAuditBridge.wrap(llm_client, container) 

222 container.singleton(LLMClientProtocol, wrapped_client) 

223 container.singleton("llm", wrapped_client) 

224 self._llm_client = wrapped_client 

225 llm_client = wrapped_client # subsequent wraps build on this 

226 logger.info("Registered LLMAuditBridge (audit_calls=True)") 

227 

228 # Cache wrap sits outermost: cache hits short-circuit before audit 

229 # fires (no real LLM call happened, so no audit entry). LXF-003 wiring 

230 # — actually invokes build_llm_cache_key in production paths. 

231 if self.config.enable_cache: 

232 cache_wrapped = LLMCacheWrapper.wrap( 

233 llm_client, 

234 cache=llm_cache, 

235 provider=self.config.provider, 

236 model=self.config.model, 

237 model_revision=getattr(self.config, "model_revision", None), 

238 ttl_seconds=getattr(self.config, "cache_ttl", None), 

239 ) 

240 container.singleton(LLMClientProtocol, cache_wrapped) 

241 container.singleton("llm", cache_wrapped) 

242 self._llm_client = cache_wrapped 

243 logger.info("Registered LLMCacheWrapper (cache hits skip audit)") 

244 

245 # Register default TokenCounterProtocol 

246 default_counter = token_registry.for_model(self.config.model or "gpt-3.5-turbo") 

247 container.singleton(TokenCounterProtocol, default_counter) 

248 

249 # Register pricing manager + cost estimator when pricing is configured 

250 if self.config.pricing and self.config.pricing.enabled: 

251 from lexigram.ai.llm.pricing.estimator import PricingCostEstimator 

252 from lexigram.ai.llm.pricing.manager import PricingManager 

253 from lexigram.contracts.ai.llm import CostEstimatorProtocol 

254 

255 pricing_manager = PricingManager( 

256 sources=self.config.pricing.build_sources(), 

257 cache_ttl=self.config.pricing.cache_ttl, 

258 enable_fuzzy_match=self.config.pricing.enable_fuzzy_match, 

259 ) 

260 container.singleton(PricingManager, pricing_manager) 

261 container.singleton( 

262 CostEstimatorProtocol, 

263 PricingCostEstimator( 

264 {}, 

265 enable_fuzzy_match=self.config.pricing.enable_fuzzy_match, 

266 ), 

267 ) 

268 logger.info( 

269 "Registered pricing manager and cost estimator", 

270 sources=[s.source_name for s in pricing_manager.sources], 

271 ) 

272 

273 # Register InstructorExtractor if instructor is available 

274 if self._instructor_available(): 

275 from lexigram.ai.llm.extraction.extractor import InstructorExtractor 

276 

277 container.singleton(InstructorExtractor, InstructorExtractor) 

278 logger.info("Registered InstructorExtractor") 

279 

280 logger.info("LLM services registered") 

281 

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

283 """Boot the LLM provider — validates API key presence and format. 

284 

285 Args: 

286 container: The DI container resolver. 

287 """ 

288 from lexigram.contracts.ai.types import ModelProvider 

289 

290 # Providers that require an API key 

291 _REQUIRES_KEY: set[str] = { 

292 ModelProvider.OPENAI, 

293 ModelProvider.ANTHROPIC, 

294 ModelProvider.COHERE, 

295 ModelProvider.GROQ, 

296 ModelProvider.MISTRAL, 

297 ModelProvider.OPENROUTER, 

298 ModelProvider.GEMINI, 

299 ModelProvider.CLOUDFLARE, 

300 ModelProvider.AZURE_OPENAI, 

301 ModelProvider.DEEPSEEK, 

302 ModelProvider.TOGETHER, 

303 ModelProvider.FIREWORKS, 

304 } 

305 

306 provider = str(self.config.provider) 

307 if provider in _REQUIRES_KEY: 

308 if not self.config.api_key: 

309 logger.warning( 

310 "llm_provider_boot_warning", 

311 provider=provider, 

312 reason="API key is not set; requests will likely fail with AuthenticationError", 

313 ) 

314 else: 

315 key_val = self.config.api_key.get_secret_value() 

316 if len(key_val) < 8: 

317 logger.warning( 

318 "llm_provider_boot_warning", 

319 provider=provider, 

320 reason="API key appears too short; verify the key is correct", 

321 ) 

322 

323 # Warm the cost estimator pricing snapshot from configured sources 

324 if self.config.pricing and self.config.pricing.enabled: 

325 try: 

326 from lexigram.ai.llm.pricing.estimator import PricingCostEstimator 

327 from lexigram.ai.llm.pricing.manager import PricingManager 

328 from lexigram.contracts.ai.llm import CostEstimatorProtocol 

329 

330 manager = await container.resolve(PricingManager) 

331 estimator = await container.resolve(CostEstimatorProtocol) 

332 await cast("PricingCostEstimator", estimator).warm(manager) 

333 except ( 

334 OSError, 

335 ValueError, 

336 TypeError, 

337 LookupError, 

338 RuntimeError, 

339 LexigramError, 

340 ModuleVisibilityError, 

341 UnresolvableDependencyError, 

342 ) as e: 

343 logger.warning( 

344 "pricing_preload_failed", 

345 error=str(e), 

346 reason="cost estimates will return 0.0 until sources are reachable", 

347 ) 

348 

349 # Optional: register this client in the provider registry 

350 if self._llm_client is not None: 

351 try: 

352 from lexigram.contracts.ai.providers import ProviderRegistryProtocol 

353 

354 registry = await container.resolve(ProviderRegistryProtocol) 

355 await registry.register_provider( 

356 name=str(self.config.provider), 

357 client=self._llm_client, 

358 models=[], 

359 ) 

360 logger.debug( 

361 "llm_registered_in_provider_registry", 

362 provider=self.config.provider, 

363 ) 

364 except ( 

365 LookupError, 

366 RuntimeError, 

367 AttributeError, 

368 ImportError, 

369 ModuleVisibilityError, 

370 UnresolvableDependencyError, 

371 ): 

372 logger.debug("llm_provider_registry_not_available") 

373 

374 async def shutdown(self) -> None: 

375 """Close client connections on application shutdown.""" 

376 if self._llm_client and hasattr(self._llm_client, "close"): 

377 try: 

378 await self._llm_client.close() 

379 except (ConnectionError, TimeoutError, OSError) as exc: 

380 logger.warning("Error closing LLM client", error=str(exc)) 

381 self._llm_client = None 

382 logger.info("LLM provider shutdown complete") 

383 

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

385 """Return basic health information for the registered LLM client.""" 

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

387 

388 return HealthCheckResult( 

389 component=self.name, 

390 status=HealthStatus.HEALTHY 

391 if self._llm_client is not None 

392 else HealthStatus.DEGRADED, 

393 details={ 

394 "provider": str(self.config.provider), 

395 "model": self.config.model, 

396 }, 

397 )