Coverage for src / lexigram / ai / relay / gateway / channels.py: 99%
70 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-08-08 23:08 +0800
1"""Deterministic channel selection for the relay gateway.
3Selection is a pure function of the static channel table, the runtime
4override table, and the routing intent: source format, model alias,
5streaming flag, requested capability flags, and an optional preferred
6channel. The registry never inspects request payloads to make an
7undocumented routing decision.
9The runtime override table is written by the operator controls surface
10(``set_runtime_enabled``); it can only take a configured channel out of
11service (drain) or restore it — it can never enable a channel whose
12config ``enabled`` flag is false.
13"""
15from __future__ import annotations
17from collections.abc import Callable
18import random
20from lexigram.ai.relay.gateway.config import RelayGatewayConfig
21from lexigram.contracts.ai.relay import RelayChannel, RelayFormat, RelayGatewayError
22from lexigram.contracts.core.result import Err, Ok, Result
24__all__ = ["RelayChannelRegistry"]
27class RelayChannelRegistry:
28 """Selects a ``RelayChannel`` deterministically from a static config.
30 By default selection is fully deterministic; when the config enables
31 ``"weighted"`` load balancing, ties among an already-tied top tier
32 (equal precedence, eligible, same priority) are broken by
33 weighted-random pick instead of ascending name.
35 Note:
36 "Healthy channel" in the plan is interpreted as the channel's
37 ``enabled`` flag combined with the live runtime override table;
38 drained channels are invisible to ``select`` until they are
39 restored at runtime.
41 Attributes:
42 _channels: The immutable channel table from ``RelayGatewayConfig``.
43 _runtime_enabled: Operator overrides applied by the controls
44 service; empty equals "all channels as configured".
45 """
47 def __init__(
48 self,
49 config: RelayGatewayConfig,
50 *,
51 random_source: Callable[[int], int] | None = None,
52 ) -> None:
53 """Bind the registry to a static channel table.
55 Args:
56 config: Immutable gateway configuration. Selection never
57 mutates it; the channel tuple is kept as configured.
58 random_source: Callable receiving a weight sum and returning
59 a value in ``[0, total)``; only used by the weighted
60 tie-break. Defaults to ``random.SystemRandom().randrange``;
61 tests pass a deterministic fake.
62 """
63 self._config = config
64 self._channels = config.channels
65 self._runtime_enabled: dict[str, bool] = {}
66 self._random_source = random_source or random.SystemRandom().randrange
68 @property
69 def channels(self) -> tuple[RelayChannel, ...]:
70 """The configured channel table.
72 Returns:
73 The immutable channel tuple, in configuration order.
74 """
75 return self._channels
77 def set_runtime_enabled(self, channel: str, enabled: bool) -> None:
78 """Override the eligibility of *channel* at runtime.
80 Draining a channel (``enabled=False``) hides it from selection
81 until restored; restoring removes the override entirely so the
82 config ``enabled`` flag alone decides. A config-disabled channel
83 can never be made eligible this way.
85 Args:
86 channel: Channel name to override.
87 enabled: Whether the channel should select new requests.
88 """
89 if enabled:
90 self._runtime_enabled.pop(channel, None)
91 else:
92 self._runtime_enabled[channel] = False
94 def runtime_enabled(self) -> dict[str, bool]:
95 """Return the non-default runtime overrides.
97 Returns:
98 Mapping of channel name to ``False`` set at runtime; a
99 restored channel is absent.
100 """
101 return dict(self._runtime_enabled)
103 def select(
104 self,
105 source: RelayFormat,
106 model: str,
107 stream: bool = False,
108 capabilities: frozenset[str] = frozenset(),
109 preferred: str | None = None,
110 exclude: frozenset[str] = frozenset(),
111 ) -> Result[RelayChannel, RelayGatewayError]:
112 """Pick the best channel for the routing query.
114 Eligibility is computed first (enabled by config and runtime,
115 target format differs from the source, model serves the
116 requested alias, streaming and capability constraints), then the
117 survivors are sorted: preferred channel first, then exact model
118 match, then ascending priority (lower number wins), then
119 ascending name as a stable tiebreak.
120 The preferred channel still must pass every eligibility filter;
121 otherwise it is skipped and normal ordering applies.
123 Args:
124 source: Wire format the caller supplies; channels whose
125 target format equals it would be no-op conversions and
126 are never eligible.
127 model: Requested model alias; only exact matches are
128 eligible.
129 stream: Whether the caller wants streaming. Channels that
130 declare capabilities must declare ``"stream"`` to serve
131 streaming requests; channels with no declared
132 capabilities are unconstrained.
133 capabilities: Requested capability flags; they must be a
134 subset of the channel's declared capabilities.
135 preferred: Optional channel name that ranks first when it is
136 eligible. Defaults to ``None`` (no preference).
137 exclude: Channel names to skip, e.g. for failover retries.
138 Excluded names are filtered before the other eligibility
139 filters run, so an excluded channel is never eligible and
140 cannot be treated as preferred. Defaults to empty (no
141 exclusion).
143 Returns:
144 ``Ok(channel)`` for the best eligible channel, or
145 ``Err(RelayGatewayError)`` when none is eligible. The error
146 cause is classified in fixed order: no enabled channels
147 (``CHANNEL_DISABLED``, 404), no enabled channel transforms
148 the source format (``TARGET_FORMAT_UNSUPPORTED``, 500), no
149 enabled channel satisfies the capability filters
150 (``CAPABILITY_UNAVAILABLE``, 409), otherwise the model is
151 not served (``MODEL_NOT_FOUND``, 404).
153 Note:
154 Runtime-drained channels are treated exactly like disabled
155 channels: they are invisible to selection until restored.
156 """
157 enabled = self._enabled_channels()
158 if not enabled:
159 return Err(
160 RelayGatewayError(
161 code="CHANNEL_DISABLED",
162 message="no enabled channels",
163 status_code=404,
164 request_id="",
165 )
166 )
167 enabled = [channel for channel in enabled if channel.name not in exclude]
168 transformable = [
169 channel for channel in enabled if channel.target_format != source
170 ]
171 if not transformable:
172 return Err(
173 RelayGatewayError(
174 code="TARGET_FORMAT_UNSUPPORTED",
175 message="no channel supports the requested target format",
176 status_code=500,
177 request_id="",
178 )
179 )
180 capable = [
181 channel
182 for channel in transformable
183 if self._meets_capabilities(channel, stream, capabilities)
184 ]
185 if not capable:
186 return Err(
187 RelayGatewayError(
188 code="CAPABILITY_UNAVAILABLE",
189 message="no channel provides the requested capabilities",
190 status_code=409,
191 request_id="",
192 )
193 )
194 matched = [channel for channel in capable if model in channel.models]
195 if not matched:
196 return Err(
197 RelayGatewayError(
198 code="MODEL_NOT_FOUND",
199 message=f"no channel serves model {model!r}",
200 status_code=404,
201 request_id="",
202 )
203 )
204 ordered = sorted(
205 matched,
206 key=lambda channel: (
207 channel.name != preferred,
208 model not in channel.models,
209 channel.priority,
210 channel.name,
211 ),
212 )
213 if self._config.load_balancing == "weighted" and len(ordered) > 1:
214 top = ordered[0]
215 top_key = (top.name != preferred, model not in top.models, top.priority)
216 tier = [
217 channel
218 for channel in ordered
219 if (
220 channel.name != preferred,
221 model not in channel.models,
222 channel.priority,
223 )
224 == top_key
225 ]
226 if len(tier) > 1:
227 return Ok(self._pick_weighted(tier))
228 return Ok(ordered[0])
230 def select_for_endpoint(
231 self,
232 kind: str,
233 model: str,
234 *,
235 exclude: frozenset[str] = frozenset(),
236 ) -> Result[RelayChannel, RelayGatewayError]:
237 """Pick the best channel serving an endpoint kind (e.g. ``"embeddings"``).
239 Passthrough entry point: eligibility is limited to channels
240 declaring *kind* in ``endpoint_kinds`` (empty means chat-only,
241 never eligible here), then survivors are sorted by ascending
242 priority (lower number wins) and ascending name as a stable
243 tiebreak. Model aliases and the ``enabled``/runtime-disabled
244 filters behave exactly like ``select``. A chat-only channel is
245 untouched by this method.
247 Args:
248 kind: Endpoint kind the caller wants (e.g. ``"embeddings"``);
249 only channels declaring it are eligible.
250 model: Requested model alias; only exact matches are
251 eligible.
252 exclude: Channel names to skip, e.g. for failover retries.
253 Defaults to empty (no exclusion).
255 Returns:
256 ``Ok(channel)`` for the best eligible channel, or
257 ``Err(RelayGatewayError)`` when none is eligible: no enabled
258 channels (``CHANNEL_DISABLED``, 404), otherwise, no channel
259 serves the kind or model (``MODEL_NOT_FOUND``, 404).
260 """
261 enabled = self._enabled_channels()
262 if not enabled:
263 return Err(
264 RelayGatewayError(
265 code="CHANNEL_DISABLED",
266 message="no enabled channels",
267 status_code=404,
268 request_id="",
269 )
270 )
271 serving = [
272 channel
273 for channel in enabled
274 if kind in channel.endpoint_kinds and channel.name not in exclude
275 ]
276 matched = [channel for channel in serving if model in channel.models]
277 if not matched:
278 return Err(
279 RelayGatewayError(
280 code="MODEL_NOT_FOUND",
281 message=f"no channel serves endpoint {kind!r} for model {model!r}",
282 status_code=404,
283 request_id="",
284 )
285 )
286 ordered = sorted(
287 matched,
288 key=lambda channel: (channel.priority, channel.name),
289 )
290 return Ok(ordered[0])
292 def _enabled_channels(self) -> list[RelayChannel]:
293 """Return channels enabled by both config and the runtime overrides."""
294 return [
295 channel
296 for channel in self._channels
297 if channel.enabled and self._runtime_enabled.get(channel.name, True)
298 ]
300 def _pick_weighted(self, tier: list[RelayChannel]) -> RelayChannel:
301 """Pick one channel from an already-tied tier by cumulative weight.
303 Channels with ``weight=0`` are excluded unless the whole tier is
304 made of them; the walk is driven by ``self._random_source`` on
305 the total weight so the pick is a pure function of the injected
306 source (low values pick early channels, values near the sum pick
307 late ones).
308 """
309 participants = [channel for channel in tier if channel.weight > 0] or tier
310 total = sum(channel.weight for channel in participants) or 1
311 roll = self._random_source(total)
312 for channel in participants:
313 roll -= channel.weight
314 if roll < 0:
315 return channel
316 return participants[-1]
318 @staticmethod
319 def _meets_capabilities(
320 channel: RelayChannel, stream: bool, capabilities: frozenset[str]
321 ) -> bool:
322 """Check whether *channel* satisfies the streaming and capability filters."""
323 if capabilities and not capabilities <= channel.capabilities:
324 return False
325 return not (
326 stream and channel.capabilities and "stream" not in channel.capabilities
327 )