1"""Parallel Race Routing Strategy."""
2
3from __future__ import annotations
4
5import asyncio
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 ParallelRaceStrategy:
26 """Fire requests to all eligible providers simultaneously.
27
28 Returns the first successful result. Remaining in-flight tasks are
29 cancelled to conserve resources and avoid unnecessary billing.
30 """
31
32 async def execute(
33 self,
34 *,
35 providers: list[ProviderConfig],
36 clients: dict[str, LLMClientProtocol],
37 quota: QuotaBackendProtocol,
38 config: LLMConfig,
39 messages: list[Any],
40 kwargs: dict[str, Any],
41 ) -> tuple[InferenceResult | None, list[str], int]:
42 temperature, max_tokens = _gen_defaults(config, kwargs)
43 providers_tried: list[str] = []
44
45 # Build list of eligible (entry config, client) pairs.
46 candidates: list[tuple[ProviderConfig, LLMClientProtocol]] = []
47 for pcfg in providers:
48 if not pcfg.enabled:
49 continue
50 client = clients.get(pcfg.key)
51 if client is None:
52 continue
53 providers_tried.append(pcfg.name)
54 if await quota.is_exhausted(pcfg.key):
55 continue
56 candidates.append((pcfg, client))
57
58 if not candidates:
59 return None, providers_tried, 0
60
61 # Fire all candidates concurrently.
62 tasks: dict[asyncio.Task[InferenceResult], ProviderConfig] = {}
63 for pcfg, client in candidates:
64 task = asyncio.create_task(
65 _attempt_provider(
66 client=client,
67 provider_name=pcfg.name,
68 model=pcfg.model,
69 messages=messages,
70 temperature=temperature,
71 max_tokens=max_tokens,
72 ),
73 )
74 tasks[task] = pcfg
75
76 total_attempts = len(tasks)
77 winner: InferenceResult | None = None
78 pending = set(tasks.keys())
79
80 try:
81 while pending:
82 done, pending = await asyncio.wait(
83 pending,
84 return_when=asyncio.FIRST_COMPLETED,
85 )
86 for task in done:
87 pcfg = tasks[task]
88 if task.exception() is not None:
89 exc = task.exception()
90 if isinstance(exc, AIError):
91 await _handle_free_failure(
92 exc=exc,
93 provider_key=pcfg.key,
94 model=pcfg.model,
95 quota=quota,
96 cooldown_seconds=config.quota.cooldown_seconds,
97 )
98 continue
99 # First success wins.
100 winner = task.result()
101 await quota.increment(pcfg.key)
102 logger.info("parallel_race: winner provider=%s", pcfg.name)
103 # Cancel remaining tasks.
104 for p in pending:
105 p.cancel()
106 break
107 finally:
108 # Ensure no dangling tasks.
109 for t in tasks:
110 if not t.done():
111 t.cancel()
112
113 return winner, providers_tried, total_attempts