1"""Latency-Optimized Routing Strategy."""
2
3from __future__ import annotations
4
5import collections
6from typing import Any
7
8from lexigram.ai.llm.routing.config import LLMConfig, ProviderConfig
9from lexigram.ai.llm.routing.strategies.base import (
10 _attempt_provider,
11 _gen_defaults,
12 _handle_free_failure,
13)
14from lexigram.ai.llm.routing.types import InferenceResult
15from lexigram.ai.llm.types import AIError
16from lexigram.contracts.ai import LLMClientProtocol
17from lexigram.contracts.ai.routing import QuotaBackendProtocol
18from lexigram.logging import (
19 get_logger,
20)
21
22logger = get_logger(__name__)
23
24
25class LatencyOptimizedStrategy:
26 """Route to the provider with the lowest recent average latency.
27
28 Maintains a rolling window of latency samples per provider.
29 Providers with no samples yet are tried first (exploration) to
30 bootstrap the statistics.
31
32 When ``skip_unhealthy`` is enabled, providers whose last
33 ``health_check()`` returned ``UNHEALTHY`` are skipped.
34
35 Args:
36 window_size: Number of recent samples to keep per provider.
37 skip_unhealthy: Run ``health_check()`` and skip unhealthy providers.
38 health_timeout: Timeout in seconds for each health check call.
39 """
40
41 def __init__(
42 self,
43 window_size: int = 20,
44 skip_unhealthy: bool = True,
45 health_timeout: float = 3.0,
46 ) -> None:
47 self._window_size = window_size
48 self._skip_unhealthy = skip_unhealthy
49 self._health_timeout = health_timeout
50 self._samples: dict[str, collections.deque[float]] = {}
51
52 def record_latency(self, provider: str, latency_ms: float) -> None:
53 """Manually record a latency sample."""
54 dq = self._samples.setdefault(
55 provider,
56 collections.deque(maxlen=self._window_size),
57 )
58 dq.append(latency_ms)
59
60 def _avg_latency(self, provider: str) -> float:
61 """Return average latency; -1.0 if no data (explore first)."""
62 dq = self._samples.get(provider)
63 if not dq:
64 return -1.0
65 return sum(dq) / len(dq)
66
67 async def execute(
68 self,
69 *,
70 providers: list[ProviderConfig],
71 clients: dict[str, LLMClientProtocol],
72 quota: QuotaBackendProtocol,
73 config: LLMConfig,
74 messages: list[Any],
75 kwargs: dict[str, Any],
76 ) -> tuple[InferenceResult | None, list[str], int]:
77 temperature, max_tokens = _gen_defaults(config, kwargs)
78
79 # Partition providers: unknown latency first (explore), then sort by avg.
80 eligible: list[ProviderConfig] = [
81 p for p in providers if p.enabled and clients.get(p.key) is not None
82 ]
83
84 # Skip unhealthy providers when enabled.
85 if self._skip_unhealthy:
86 healthy: list[ProviderConfig] = []
87 for p in eligible:
88 client = clients[p.key]
89 if hasattr(client, "health_check"):
90 try:
91 hc = await client.health_check(timeout=self._health_timeout)
92 if hasattr(hc, "status") and str(hc.status) == "unhealthy":
93 logger.info(
94 "latency_optimized: skipping unhealthy provider=%s",
95 p.name,
96 )
97 continue
98 except (OSError, ConnectionError, TimeoutError, RuntimeError):
99 logger.debug(
100 "latency_optimized: health_check failed for %s, including anyway",
101 p.name,
102 )
103 healthy.append(p)
104 # Fall back to full list if all are unhealthy.
105 if healthy:
106 eligible = healthy
107
108 unknown = [p for p in eligible if self._avg_latency(p.key) < 0]
109 known = [p for p in eligible if self._avg_latency(p.key) >= 0]
110 known.sort(key=lambda p: self._avg_latency(p.key))
111
112 ordered = unknown + known
113
114 providers_tried: list[str] = []
115 total_attempts = 0
116
117 for pcfg in ordered:
118 client = clients[pcfg.key]
119 providers_tried.append(pcfg.name)
120
121 if await quota.is_exhausted(pcfg.key):
122 continue
123
124 model = pcfg.model
125 total_attempts += 1
126 try:
127 result = await _attempt_provider(
128 client=client,
129 provider_name=pcfg.name,
130 model=model,
131 messages=messages,
132 temperature=temperature,
133 max_tokens=max_tokens,
134 )
135 await quota.increment(pcfg.key)
136 self.record_latency(pcfg.key, result.latency_ms)
137 logger.info(
138 "latency_optimized: success provider=%s model=%s "
139 "latency_ms=%.0f avg_latency=%.0f",
140 pcfg.name,
141 model,
142 result.latency_ms,
143 self._avg_latency(pcfg.key),
144 )
145 return result, providers_tried, total_attempts
146 except AIError as exc:
147 await _handle_free_failure(
148 exc=exc,
149 provider_key=pcfg.key,
150 model=model,
151 quota=quota,
152 cooldown_seconds=config.quota.cooldown_seconds,
153 )
154
155 return None, providers_tried, total_attempts