1"""Deterministic channel selection for the relay gateway.
2
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.
8
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"""
14
15from __future__ import annotations
16
17from collections.abc import Callable
18import random
19
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
23
24__all__ = ["RelayChannelRegistry"]
25
26
27class RelayChannelRegistry:
28 """Selects a ``RelayChannel`` deterministically from a static config.
29
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.
34
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.
40
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 """
46
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.
54
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
67
68 @property
69 def channels(self) -> tuple[RelayChannel, ...]:
70 """The configured channel table.
71
72 Returns:
73 The immutable channel tuple, in configuration order.
74 """
75 return self._channels
76
77 def reload(self, channels: tuple[RelayChannel, ...]) -> None:
78 """Replace the channel table (boot reconcile from a durable store).
79
80 Runtime overrides are preserved for channels that remain in the
81 new table and dropped for channels that were removed, so a boot
82 reconcile never resurrects a drained channel and never carries
83 overrides for channels that no longer exist.
84
85 Args:
86 channels: The new channel tuple, in selection order.
87 """
88 remaining = {channel.name for channel in channels}
89 self._runtime_enabled = {
90 name: enabled
91 for name, enabled in self._runtime_enabled.items()
92 if name in remaining
93 }
94 self._channels = channels
95
96 def set_runtime_enabled(self, channel: str, enabled: bool) -> None:
97 """Override the eligibility of *channel* at runtime.
98
99 Draining a channel (``enabled=False``) hides it from selection
100 until restored; restoring removes the override entirely so the
101 config ``enabled`` flag alone decides. A config-disabled channel
102 can never be made eligible this way.
103
104 Args:
105 channel: Channel name to override.
106 enabled: Whether the channel should select new requests.
107 """
108 if enabled:
109 self._runtime_enabled.pop(channel, None)
110 else:
111 self._runtime_enabled[channel] = False
112
113 def runtime_enabled(self) -> dict[str, bool]:
114 """Return the non-default runtime overrides.
115
116 Returns:
117 Mapping of channel name to ``False`` set at runtime; a
118 restored channel is absent.
119 """
120 return dict(self._runtime_enabled)
121
122 def select(
123 self,
124 source: RelayFormat,
125 model: str,
126 stream: bool = False,
127 capabilities: frozenset[str] = frozenset(),
128 preferred: str | None = None,
129 exclude: frozenset[str] = frozenset(),
130 ) -> Result[RelayChannel, RelayGatewayError]:
131 """Pick the best channel for the routing query.
132
133 Eligibility is computed first (enabled by config and runtime,
134 target format differs from the source, model serves the
135 requested alias, streaming and capability constraints), then the
136 survivors are sorted: preferred channel first, then exact model
137 match, then ascending priority (lower number wins), then
138 ascending name as a stable tiebreak.
139 The preferred channel still must pass every eligibility filter;
140 otherwise it is skipped and normal ordering applies.
141
142 Args:
143 source: Wire format the caller supplies; channels whose
144 target format equals it would be no-op conversions and
145 are never eligible.
146 model: Requested model alias; only exact matches are
147 eligible.
148 stream: Whether the caller wants streaming. Channels that
149 declare capabilities must declare ``"stream"`` to serve
150 streaming requests; channels with no declared
151 capabilities are unconstrained.
152 capabilities: Requested capability flags; they must be a
153 subset of the channel's declared capabilities.
154 preferred: Optional channel name that ranks first when it is
155 eligible. Defaults to ``None`` (no preference).
156 exclude: Channel names to skip, e.g. for failover retries.
157 Excluded names are filtered before the other eligibility
158 filters run, so an excluded channel is never eligible and
159 cannot be treated as preferred. Defaults to empty (no
160 exclusion).
161
162 Returns:
163 ``Ok(channel)`` for the best eligible channel, or
164 ``Err(RelayGatewayError)`` when none is eligible. The error
165 cause is classified in fixed order: no enabled channels
166 (``CHANNEL_DISABLED``, 404), no enabled channel transforms
167 the source format (``TARGET_FORMAT_UNSUPPORTED``, 500), no
168 enabled channel satisfies the capability filters
169 (``CAPABILITY_UNAVAILABLE``, 409), otherwise the model is
170 not served (``MODEL_NOT_FOUND``, 404).
171
172 Note:
173 Runtime-drained channels are treated exactly like disabled
174 channels: they are invisible to selection until restored.
175 """
176 enabled = self._enabled_channels()
177 if not enabled:
178 return Err(
179 RelayGatewayError(
180 code="CHANNEL_DISABLED",
181 message="no enabled channels",
182 status_code=404,
183 request_id="",
184 )
185 )
186 enabled = [channel for channel in enabled if channel.name not in exclude]
187 transformable = [
188 channel for channel in enabled if channel.target_format != source
189 ]
190 if not transformable:
191 return Err(
192 RelayGatewayError(
193 code="TARGET_FORMAT_UNSUPPORTED",
194 message="no channel supports the requested target format",
195 status_code=500,
196 request_id="",
197 )
198 )
199 capable = [
200 channel
201 for channel in transformable
202 if self._meets_capabilities(channel, stream, capabilities)
203 ]
204 if not capable:
205 return Err(
206 RelayGatewayError(
207 code="CAPABILITY_UNAVAILABLE",
208 message="no channel provides the requested capabilities",
209 status_code=409,
210 request_id="",
211 )
212 )
213 matched = [channel for channel in capable if model in channel.models]
214 if not matched:
215 return Err(
216 RelayGatewayError(
217 code="MODEL_NOT_FOUND",
218 message=f"no channel serves model {model!r}",
219 status_code=404,
220 request_id="",
221 )
222 )
223 ordered = sorted(
224 matched,
225 key=lambda channel: (
226 channel.name != preferred,
227 model not in channel.models,
228 channel.priority,
229 channel.name,
230 ),
231 )
232 if self._config.load_balancing == "weighted" and len(ordered) > 1:
233 top = ordered[0]
234 top_key = (top.name != preferred, model not in top.models, top.priority)
235 tier = [
236 channel
237 for channel in ordered
238 if (
239 channel.name != preferred,
240 model not in channel.models,
241 channel.priority,
242 )
243 == top_key
244 ]
245 if len(tier) > 1:
246 return Ok(self._pick_weighted(tier))
247 return Ok(ordered[0])
248
249 def select_for_endpoint(
250 self,
251 kind: str,
252 model: str,
253 *,
254 exclude: frozenset[str] = frozenset(),
255 ) -> Result[RelayChannel, RelayGatewayError]:
256 """Pick the best channel serving an endpoint kind (e.g. ``"embeddings"``).
257
258 Passthrough entry point: eligibility is limited to channels
259 declaring *kind* in ``endpoint_kinds`` (empty means chat-only,
260 never eligible here), then survivors are sorted by ascending
261 priority (lower number wins) and ascending name as a stable
262 tiebreak. Model aliases and the ``enabled``/runtime-disabled
263 filters behave exactly like ``select``. A chat-only channel is
264 untouched by this method.
265
266 Args:
267 kind: Endpoint kind the caller wants (e.g. ``"embeddings"``);
268 only channels declaring it are eligible.
269 model: Requested model alias; only exact matches are
270 eligible.
271 exclude: Channel names to skip, e.g. for failover retries.
272 Defaults to empty (no exclusion).
273
274 Returns:
275 ``Ok(channel)`` for the best eligible channel, or
276 ``Err(RelayGatewayError)`` when none is eligible: no enabled
277 channels (``CHANNEL_DISABLED``, 404), otherwise, no channel
278 serves the kind or model (``MODEL_NOT_FOUND``, 404).
279 """
280 enabled = self._enabled_channels()
281 if not enabled:
282 return Err(
283 RelayGatewayError(
284 code="CHANNEL_DISABLED",
285 message="no enabled channels",
286 status_code=404,
287 request_id="",
288 )
289 )
290 serving = [
291 channel
292 for channel in enabled
293 if kind in channel.endpoint_kinds and channel.name not in exclude
294 ]
295 matched = [channel for channel in serving if model in channel.models]
296 if not matched:
297 return Err(
298 RelayGatewayError(
299 code="MODEL_NOT_FOUND",
300 message=f"no channel serves endpoint {kind!r} for model {model!r}",
301 status_code=404,
302 request_id="",
303 )
304 )
305 ordered = sorted(
306 matched,
307 key=lambda channel: (channel.priority, channel.name),
308 )
309 return Ok(ordered[0])
310
311 def _enabled_channels(self) -> list[RelayChannel]:
312 """Return channels enabled by both config and the runtime overrides."""
313 return [
314 channel
315 for channel in self._channels
316 if channel.enabled and self._runtime_enabled.get(channel.name, True)
317 ]
318
319 def _pick_weighted(self, tier: list[RelayChannel]) -> RelayChannel:
320 """Pick one channel from an already-tied tier by cumulative weight.
321
322 Channels with ``weight=0`` are excluded unless the whole tier is
323 made of them; the walk is driven by ``self._random_source`` on
324 the total weight so the pick is a pure function of the injected
325 source (low values pick early channels, values near the sum pick
326 late ones).
327 """
328 participants = [channel for channel in tier if channel.weight > 0] or tier
329 total = sum(channel.weight for channel in participants) or 1
330 roll = self._random_source(total)
331 for channel in participants:
332 roll -= channel.weight
333 if roll < 0:
334 return channel
335 return participants[-1]
336
337 @staticmethod
338 def _meets_capabilities(
339 channel: RelayChannel, stream: bool, capabilities: frozenset[str]
340 ) -> bool:
341 """Check whether *channel* satisfies the streaming and capability filters."""
342 if capabilities and not capabilities <= channel.capabilities:
343 return False
344 return not (
345 stream and channel.capabilities and "stream" not in channel.capabilities
346 )