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

1from __future__ import annotations 

2 

3from dataclasses import dataclass 

4import hashlib 

5from typing import Any 

6 

7from lexigram.contracts.ai.llm import LLMClientProtocol 

8from lexigram.logging import ( 

9 get_logger, 

10) 

11 

12logger = get_logger(__name__) 

13 

14__all__ = ["ABSplitConfig", "ABSplitStrategy"] 

15 

16 

17@dataclass(frozen=True) 

18class ABSplitConfig: 

19 """Configuration for A/B traffic splitting between two LLM providers. 

20 

21 Attributes: 

22 control_key: Container key for the control (baseline) provider. 

23 treatment_key: Container key for the treatment (experimental) provider. 

24 treatment_percentage: 0–100, percentage of traffic to route to treatment. 

25 split_key_field: Request field used to deterministically assign traffic. 

26 """ 

27 

28 control_key: str 

29 treatment_key: str 

30 treatment_percentage: int = 10 

31 split_key_field: str = "user_id" 

32 

33 def __post_init__(self) -> None: 

34 """Validate percentage range. 

35 

36 Raises: 

37 ValueError: If treatment_percentage is not in [0, 100]. 

38 """ 

39 if not 0 <= self.treatment_percentage <= 100: 

40 raise ValueError( 

41 f"treatment_percentage must be 0-100, got {self.treatment_percentage}" 

42 ) 

43 

44 

45class ABSplitStrategy: 

46 """Routes LLM requests between control and treatment providers using deterministic hashing. 

47 

48 Uses MD5 hashing on the split key to ensure consistent assignment: 

49 the same user always gets the same variant for a given configuration. 

50 

51 Args: 

52 config: A/B split configuration. 

53 control: The control (baseline) LLM provider. 

54 treatment: The treatment (experimental) LLM provider. 

55 """ 

56 

57 def __init__( 

58 self, 

59 config: ABSplitConfig, 

60 control: LLMClientProtocol, 

61 treatment: LLMClientProtocol, 

62 ) -> None: 

63 self._config = config 

64 self._control = control 

65 self._treatment = treatment 

66 

67 def _should_use_treatment(self, request: Any) -> bool: 

68 """Deterministically assign request to control or treatment. 

69 

70 Hashes ``split_key_field`` from the request (or ``str(request)`` as 

71 fallback) to a value in [0, 99], then compares against 

72 ``treatment_percentage``. 

73 

74 Args: 

75 request: The model request (any object with a split_key_field attribute). 

76 

77 Returns: 

78 True if request should go to treatment, False for control. 

79 """ 

80 split_value = str(getattr(request, self._config.split_key_field, "")) 

81 if not split_value: 

82 return False 

83 

84 hash_bytes = hashlib.md5( # noqa: S324 

85 f"{split_value}:{self._config.treatment_key}".encode(), 

86 usedforsecurity=False, 

87 ).digest() 

88 bucket = int.from_bytes(hash_bytes[:4], "big") % 100 

89 return bucket < self._config.treatment_percentage 

90 

91 async def route(self, request: Any) -> LLMClientProtocol: 

92 """Select the appropriate provider for this request. 

93 

94 Args: 

95 request: The model request to route. 

96 

97 Returns: 

98 Either the control or treatment LLM provider. 

99 """ 

100 use_treatment = self._should_use_treatment(request) 

101 variant = "treatment" if use_treatment else "control" 

102 logger.debug("ab_split_routed", variant=variant) 

103 return self._treatment if use_treatment else self._control