Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-llm/src/lexigram/ai/llm/routing/strategies/latency.py: 25%

65 statements  

« prev     ^ index     » next       coverage.py v7.15.4, created at 2026-08-25 07:19 +0800

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