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

55 statements  

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

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