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

1"""Cost-Optimized Routing Strategy.""" 

2 

3from __future__ import annotations 

4 

5from typing import TYPE_CHECKING, Any 

6 

7from lexigram.ai.llm.routing.config import LLMConfig, ProviderConfig 

8from lexigram.ai.llm.routing.strategies.base import ( 

9 _attempt_provider, 

10 _estimate_prompt_tokens, 

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.contracts.exceptions.base import LexigramError 

19from lexigram.logging import ( 

20 get_logger, 

21) 

22 

23if TYPE_CHECKING: 

24 from lexigram.ai.llm.pricing.manager import PricingManager 

25 

26logger = get_logger(__name__) 

27 

28 

29class CostOptimizedStrategy: 

30 """Sort providers by per-token cost and try cheapest first. 

31 

32 Requires a ``PricingManager`` for cost lookups. Providers whose 

33 pricing is unavailable are pushed to the end of the queue. 

34 """ 

35 

36 def __init__(self, pricing_manager: PricingManager) -> None: 

37 self._pricing = pricing_manager 

38 

39 async def execute( 

40 self, 

41 *, 

42 providers: list[ProviderConfig], 

43 clients: dict[str, LLMClientProtocol], 

44 quota: QuotaBackendProtocol, 

45 config: LLMConfig, 

46 messages: list[Any], 

47 kwargs: dict[str, Any], 

48 ) -> tuple[InferenceResult | None, list[str], int]: 

49 temperature, max_tokens = _gen_defaults(config, kwargs) 

50 

51 # Estimate token counts for cost scoring. 

52 estimated_prompt = _estimate_prompt_tokens(messages) 

53 estimated_completion = max_tokens or 500 

54 

55 # Sort by estimated cost; unknown cost → infinity. 

56 scored: list[tuple[float, ProviderConfig]] = [] 

57 for pcfg in providers: 

58 if not pcfg.enabled: 

59 continue 

60 if clients.get(pcfg.key) is None: 

61 continue 

62 try: 

63 pricing = await self._pricing.get_pricing(pcfg.model) 

64 if pricing is None: 

65 cost = float("inf") 

66 else: 

67 cost = (estimated_prompt / 1_000_000) * pricing.prompt_per_1m + ( 

68 estimated_completion / 1_000_000 

69 ) * pricing.completion_per_1m 

70 except ( 

71 ValueError, 

72 TypeError, 

73 AttributeError, 

74 LookupError, 

75 OSError, 

76 LexigramError, 

77 ): 

78 cost = float("inf") 

79 scored.append((cost, pcfg)) 

80 

81 scored.sort(key=lambda pair: pair[0]) 

82 

83 # Delegate to sequential over the cost-sorted order. 

84 providers_tried: list[str] = [] 

85 total_attempts = 0 

86 

87 for _cost, pcfg in scored: 

88 client = clients[pcfg.key] 

89 providers_tried.append(pcfg.name) 

90 

91 if await quota.is_exhausted(pcfg.key): 

92 continue 

93 

94 model = pcfg.model 

95 total_attempts += 1 

96 try: 

97 result = await _attempt_provider( 

98 client=client, 

99 provider_name=pcfg.name, 

100 model=model, 

101 messages=messages, 

102 temperature=temperature, 

103 max_tokens=max_tokens, 

104 ) 

105 await quota.increment(pcfg.key) 

106 logger.info( 

107 "cost_optimized: success provider=%s model=%s cost_score=%.4f", 

108 pcfg.name, 

109 model, 

110 _cost, 

111 ) 

112 return result, providers_tried, total_attempts 

113 except AIError as exc: 

114 await _handle_free_failure( 

115 exc=exc, 

116 provider_key=pcfg.key, 

117 model=model, 

118 quota=quota, 

119 cooldown_seconds=config.quota.cooldown_seconds, 

120 ) 

121 

122 return None, providers_tried, total_attempts