Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-llm/src/lexigram/ai/llm/routing/router.py: 17%
128 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 07:19 +0800
1"""Multi-provider LLM router.
3Cascades inference requests across configured providers in priority order.
4Returns ``Result[InferenceLog, InferenceError]`` for explicit error handling.
6The routing *strategy* is pluggable — see :mod:`.strategies` for the
7``RoutingStrategyProtocol`` protocol and built-in implementations.
8"""
10from __future__ import annotations
12from typing import TYPE_CHECKING, Any
14from lexigram.ai.llm.routing.strategies import (
15 CostOptimizedStrategy,
16 LatencyOptimizedStrategy,
17 ParallelRaceStrategy,
18 RoutingStrategyProtocol,
19 SequentialCascadeStrategy,
20)
21from lexigram.ai.llm.routing.types import (
22 InferenceError,
23 InferenceLog,
24 InferenceResult,
25)
26from lexigram.contracts.web import HttpStatusError
27from lexigram.logging import (
28 get_logger,
29)
30from lexigram.result import Err, Ok, Result
32if TYPE_CHECKING:
33 from lexigram.ai.llm.routing.config import LLMConfig
34 from lexigram.ai.llm.selection.core import ModelSelector
35 from lexigram.contracts.ai import LLMClientProtocol
36 from lexigram.contracts.ai.providers import ProviderRegistryProtocol
37 from lexigram.contracts.ai.routing import (
38 InferenceLoggerProtocol,
39 QuotaBackendProtocol,
40 )
42logger = get_logger(__name__)
44__all__ = ["LLMRouter"]
47class LLMRouter:
48 """Multi-provider LLM router with pluggable strategy.
50 Delegates to a ``RoutingStrategyProtocol`` for provider selection.
51 Returns ``Result[InferenceLog, InferenceError]``.
53 Example:
54 >>> config = LLMConfig.from_env()
55 >>> router = LLMRouter(clients=clients, quota_backend=backend, ...)
56 >>> result = await router.route(messages=[{"role": "user", "content": "Hi"}])
57 >>> if result.is_ok():
58 ... print(result.unwrap().result.content)
59 """
61 def __init__(
62 self,
63 clients: dict[str, LLMClientProtocol],
64 quota_backend: QuotaBackendProtocol,
65 inference_logger: InferenceLoggerProtocol,
66 config: LLMConfig,
67 strategy: RoutingStrategyProtocol | None = None,
68 model_selector: ModelSelector | None = None,
69 provider_registry: ProviderRegistryProtocol | None = None,
70 ) -> None:
71 """Initialise the router.
73 Args:
74 clients: Mapping of provider name to its ``LLMClientProtocol``.
75 quota_backend: Quota tracking backend.
76 inference_logger: Inference log persistence backend.
77 config: Full routing configuration.
78 strategy: Routing strategy; defaults to
79 :class:`SequentialCascadeStrategy`.
80 model_selector: Optional model selector for capability-based
81 filtering. When set, ``required_capabilities`` in kwargs
82 are used to exclude providers whose models lack the
83 requested capabilities.
84 provider_registry: Optional provider registry from
85 ``lexigram-ai-providers``. When set, additional clients
86 registered in the provider registry are merged into the
87 routing pool.
88 """
89 self._clients = clients
90 self._quota = quota_backend
91 self._logger = inference_logger
92 self._config = config
93 self._strategy: RoutingStrategyProtocol = (
94 strategy or self._strategy_from_config(config)
95 )
96 self._model_selector = model_selector
97 self._provider_registry = provider_registry
98 self._closed = False
100 # Merge provider registry clients into routing pool
101 if self._provider_registry:
102 self._merge_registry_clients()
104 @property
105 def strategy(self) -> RoutingStrategyProtocol:
106 """Return the active routing strategy."""
107 return self._strategy
109 @strategy.setter
110 def strategy(self, value: RoutingStrategyProtocol) -> None:
111 """Hot-swap the routing strategy at runtime."""
112 self._strategy = value
114 def _merge_registry_clients(self) -> None:
115 """Merge clients from the provider registry into the routing pool.
117 Only adds clients for providers not already in the pool, avoiding
118 duplicates. This is called once at construction time.
119 """
120 if self._provider_registry is None:
121 return
122 try:
123 for name in self._provider_registry.list_providers():
124 if name not in self._clients:
125 # get_client is async but we can't await here; defer to route()
126 logger.debug(
127 "router: provider_registry has additional provider %s",
128 name,
129 )
130 except (RuntimeError, TypeError, AttributeError) as e:
131 logger.warning("router: failed to merge registry clients", error=str(e))
133 async def route(
134 self,
135 messages: list[Any],
136 **kwargs: Any,
137 ) -> Result[InferenceLog, InferenceError]:
138 """Route a completion request across configured providers.
140 Args:
141 messages: OpenAI-compatible message list.
142 **kwargs: Generation overrides forwarded to client (temperature,
143 max_tokens, model, etc.).
145 Returns:
146 ``Ok(InferenceLog)`` on success, ``Err(InferenceError)`` when
147 every provider (including paid fallback) is exhausted or fails.
148 """
149 result, providers_tried, total_attempts = await self._strategy.execute(
150 providers=self._filter_by_capabilities(
151 self._config.providers,
152 kwargs,
153 ),
154 clients=self._clients,
155 quota=self._quota,
156 config=self._config,
157 messages=messages,
158 kwargs=kwargs,
159 )
161 if result is not None:
162 log = InferenceLog(
163 result=result,
164 providers_tried=providers_tried,
165 total_attempts=total_attempts,
166 )
167 await self._logger.log(log)
168 return Ok(log)
170 terminal_error = InferenceError(
171 message="All LLM providers exhausted or unavailable",
172 providers_tried=providers_tried,
173 )
174 log = InferenceLog(
175 error=terminal_error,
176 providers_tried=providers_tried,
177 total_attempts=total_attempts,
178 )
179 await self._logger.log(log)
180 return Err(terminal_error)
182 async def health_probe(self) -> Result[InferenceLog, InferenceError]:
183 """Probe configured providers without performing normal inference.
185 Returns:
186 ``Ok(InferenceLog)`` for the first healthy enabled provider.
187 The returned log is synthetic and is **not** persisted to the
188 inference logger. ``Err(InferenceError)`` is returned when no
189 healthy provider can be found.
190 """
191 providers_tried: list[str] = []
192 total_attempts = 0
194 for provider_config in self._config.providers:
195 if not provider_config.enabled:
196 continue
198 client = self._clients.get(provider_config.key)
199 if client is None:
200 continue
202 health_check = getattr(client, "health_check", None)
203 if not callable(health_check):
204 continue
206 providers_tried.append(provider_config.name)
207 total_attempts += 1
208 try:
209 health_result = await health_check(timeout=5.0)
210 except (
211 OSError,
212 ConnectionError,
213 TimeoutError,
214 RuntimeError,
215 HttpStatusError,
216 ) as exc:
217 logger.debug(
218 "llm.router: health probe provider check failed",
219 provider=provider_config.name,
220 error=str(exc),
221 )
222 continue
224 try:
225 is_healthy = health_result.is_healthy()
226 except (AttributeError, TypeError) as exc:
227 logger.debug(
228 "llm.router: health probe provider check failed",
229 provider=provider_config.name,
230 error=str(exc),
231 )
232 continue
234 if is_healthy:
235 try:
236 duration_ms = health_result.duration_ms
237 except (AttributeError, TypeError) as exc:
238 logger.debug(
239 "llm.router: health probe provider check failed",
240 provider=provider_config.name,
241 error=str(exc),
242 )
243 continue
245 return Ok(
246 InferenceLog(
247 result=InferenceResult(
248 provider=provider_config.name,
249 model=provider_config.model,
250 content="health_probe",
251 attempt=total_attempts,
252 prompt_tokens=0,
253 completion_tokens=0,
254 latency_ms=duration_ms,
255 ),
256 providers_tried=providers_tried,
257 total_attempts=total_attempts,
258 context={"health_probe": True},
259 )
260 )
262 try:
263 status = str(health_result.status)
264 message = health_result.message
265 error = health_result.error
266 except (AttributeError, TypeError) as exc:
267 logger.debug(
268 "llm.router: health probe provider check failed",
269 provider=provider_config.name,
270 error=str(exc),
271 )
272 continue
274 logger.debug(
275 "llm.router: health probe provider unhealthy",
276 provider=provider_config.name,
277 status=status,
278 message=message,
279 error=error,
280 )
282 return Err(
283 InferenceError(
284 message="No healthy LLM providers available for health probe",
285 providers_tried=providers_tried,
286 )
287 )
289 async def close(self) -> None:
290 """Close all registered LLM clients and release connections.
292 Iterates over every client in the provider pool and calls ``close()``
293 if the method exists. Safe to call multiple times.
294 """
295 if self._closed:
296 return
297 self._closed = True
298 for name, client in self._clients.items():
299 closer = getattr(client, "close", None)
300 if closer is not None:
301 try:
302 await closer()
303 logger.debug("llm.router: closed client %s", name)
304 except Exception as e:
305 logger.warning(
306 "llm.router: error closing client",
307 client=name,
308 error=str(e),
309 )
311 def _filter_by_capabilities(
312 self,
313 providers: list[Any],
314 kwargs: dict[str, Any],
315 ) -> list[Any]:
316 """Filter providers whose models lack required capabilities.
318 When ``required_capabilities`` is present in *kwargs* and a
319 ``ModelSelector`` is configured, only providers whose primary
320 model satisfies **all** requested capabilities are retained.
321 Providers whose model is unknown to the selector are kept
322 (fail-open) so the strategy can still attempt them.
324 Returns:
325 Filtered provider list (original list when no filtering applies).
326 """
327 required: list[str] | None = kwargs.get("required_capabilities")
328 if not required or self._model_selector is None:
329 return providers
331 capable_models = set(
332 self._model_selector._filter_by_capabilities(required),
333 )
335 filtered = []
336 for pcfg in providers:
337 model = pcfg.model
338 # Keep providers whose model is in the capable set or unknown
339 # to the selector (fail-open avoids silently dropping custom models).
340 if (
341 model in capable_models
342 or model not in self._model_selector.model_capabilities
343 ):
344 filtered.append(pcfg)
345 else:
346 logger.debug(
347 "router: skipping provider %s (model=%s) — lacks capabilities %s",
348 pcfg.name,
349 model,
350 required,
351 )
353 # If filtering removed everything, fall back to original list to
354 # avoid a guaranteed "all providers exhausted" error.
355 if not filtered:
356 logger.warning(
357 "router: no providers match required_capabilities %s; "
358 "falling back to full provider list",
359 required,
360 )
361 return providers
363 return filtered
365 @staticmethod
366 def _strategy_from_config(config: LLMConfig) -> RoutingStrategyProtocol:
367 """Build a :class:`RoutingStrategyProtocol` from the ``strategy`` config field.
369 Falls back to :class:`SequentialCascadeStrategy` for unknown values.
370 ``CostOptimizedStrategy`` requires a ``PricingManager``; when none
371 is available it falls back to sequential with a warning.
372 """
373 name = getattr(config, "strategy", "sequential")
374 if name == "parallel_race":
375 return ParallelRaceStrategy()
376 if name == "cost_optimized":
377 try:
378 from lexigram.ai.llm.pricing import PricingManager
380 return CostOptimizedStrategy(PricingManager.from_defaults())
381 except (ImportError, RuntimeError, AttributeError, ValueError):
382 logger.warning(
383 "cost_optimized strategy requested but PricingManager "
384 "unavailable; falling back to sequential",
385 )
386 return SequentialCascadeStrategy()
387 if name == "latency_optimized":
388 return LatencyOptimizedStrategy()
389 return SequentialCascadeStrategy()