Coverage for qreward / client / model_router.py: 98.11%

1"""Model-based routing for directing requests to specific proxy groups.""" 

2 

3import fnmatch (empty)

4import threading (empty)

5from dataclasses import dataclass, field (empty)

6from typing import Optional (empty)

7 

8from qreward.client.load_balancer import ( (empty)

9 LoadBalanceStrategy, 

10 RoundRobinSelector, 

11 WeightedRoundRobinSelector, 

12) 

13from qreward.client.openai import OpenAIChatProxy (empty)

14 

15 

16@dataclass (empty)

17class ProxyGroup: (empty)

18 """A group of proxies serving a specific model or model pattern. 

19 

20 Attributes: 

21 pattern: Model name or glob pattern (e.g. "gpt-4" or "gpt-*"). 

22 proxies: Mapping of key → proxy instance. 

23 insertion_order: Keys in insertion order for round-robin. 

24 healthy_keys: Set of currently healthy proxy keys. 

25 weights: Per-key weights for weighted round-robin. 

26 strategy: Load balancing strategy for this group. 

27 """ 

28 

29 pattern: str (empty)

30 proxies: dict[str, OpenAIChatProxy] = field(default_factory=dict) (empty)

31 insertion_order: list[str] = field(default_factory=list) (empty)

32 healthy_keys: set[str] = field(default_factory=set) (empty)

33 weights: dict[str, int] = field(default_factory=dict) (empty)

34 strategy: LoadBalanceStrategy = LoadBalanceStrategy.ROUND_ROBIN (empty)

35 rr_selector: RoundRobinSelector = field(default_factory=RoundRobinSelector) (empty)

36 wrr_selector: WeightedRoundRobinSelector = field( (empty)

37 default_factory=WeightedRoundRobinSelector 

38 ) 

39 

40 

41class ModelRouter: (empty)

42 """Routes model names to dedicated proxy groups. 

43 

44 Matching priority: exact match → glob/wildcard match → None (fallback). 

45 

46 Example: 

47 >>> router = ModelRouter() 

48 >>> router.register("gpt-4", group) 

49 >>> router.register("gpt-*", wildcard_group) 

50 >>> router.resolve("gpt-4") # exact match 

51 >>> router.resolve("gpt-3.5") # wildcard match 

52 >>> router.resolve("claude-3") # None (no match) 

53 """ 

54 

55 def __init__(self) -> None: (empty)

56 self._exact: dict[str, ProxyGroup] = {} (empty)

57 self._wildcards: list[ProxyGroup] = [] (empty)

58 self._lock = threading.Lock() (empty)

59 

60 def register( (empty)

61 self, 

62 pattern: str, 

63 proxy_map: dict[str, OpenAIChatProxy], 

64 weights: Optional[dict[str, int]] = None, 

65 strategy: LoadBalanceStrategy = LoadBalanceStrategy.ROUND_ROBIN, 

66 ) -> None: 

67 """Register a proxy group for a model pattern. 

68 

69 Args: 

70 pattern: Exact model name or glob pattern (e.g. "gpt-*"). 

71 proxy_map: Mapping of key → OpenAIChatProxy. 

72 weights: Optional per-key weights for weighted round-robin. 

73 strategy: Load balancing strategy for this group. 

74 """ 

75 group = ProxyGroup( (empty)

76 pattern=pattern, 

77 proxies=dict(proxy_map), 

78 insertion_order=list(proxy_map.keys()), 

79 healthy_keys=set(proxy_map.keys()), 

80 weights=weights or {k: 1 for k in proxy_map}, 

81 strategy=strategy, 

82 ) 

83 group.wrr_selector.update_weights(group.weights) (empty)

84 

85 with self._lock: (empty)

86 if _is_glob_pattern(pattern): (empty)

87 self._wildcards.append(group) (empty)

88 else: 

89 self._exact[pattern] = group (empty)

90 

91 def resolve(self, model: str) -> Optional[ProxyGroup]: (empty)

92 """Find the best matching proxy group for a model name. 

93 

94 Args: 

95 model: The model name to route. 

96 

97 Returns: 

98 Matching ProxyGroup or None if no route matches. 

99 """ 

100 with self._lock: (empty)

101 if model in self._exact: (empty)

102 return self._exact[model] (empty)

103 for group in self._wildcards: (empty)

104 if fnmatch.fnmatch(model, group.pattern): (empty)

105 return group (empty)

106 return None (empty)

107 

108 def select_from_group(self, group: ProxyGroup) -> OpenAIChatProxy: (empty)

109 """Select a proxy from a group using its load balancing strategy. 

110 

111 Args: 

112 group: The proxy group to select from. 

113 

114 Returns: 

115 Selected OpenAIChatProxy instance. 

116 

117 Raises: 

118 RuntimeError: When no healthy proxy is available in the group. 

119 """ 

120 selected_key: Optional[str] = None (empty)

121 

122 if group.strategy == LoadBalanceStrategy.ROUND_ROBIN: (empty)

123 selected_key = group.rr_selector.select( (empty)

124 group.insertion_order, group.healthy_keys 

125 ) 

126 elif group.strategy == LoadBalanceStrategy.WEIGHTED_ROUND_ROBIN: (empty)

127 selected_key = group.wrr_selector.select( (empty)

128 group.insertion_order, group.healthy_keys, group.weights 

129 ) 

130 elif group.strategy == LoadBalanceStrategy.LEAST_CONNECTIONS: 130 ↛ 135line 130 didn't jump to line 135 because the condition on line 130 was always true(empty)

131 selected_key = group.rr_selector.select( (empty)

132 group.insertion_order, group.healthy_keys 

133 ) 

134 

135 if selected_key is None: (empty)

136 raise RuntimeError( (empty)

137 f"No healthy proxy in group '{group.pattern}'." 

138 ) 

139 return group.proxies[selected_key] (empty)

140 

141 def mark_unhealthy(self, pattern: str, key: str) -> None: (empty)

142 """Mark a proxy as unhealthy within a specific group. 

143 

144 Args: 

145 pattern: The group pattern. 

146 key: The proxy key within the group. 

147 """ 

148 group = self._find_group(pattern) (empty)

149 if group and key in group.proxies: (empty)

150 group.healthy_keys.discard(key) (empty)

151 

152 def mark_healthy(self, pattern: str, key: str) -> None: (empty)

153 """Mark a proxy as healthy within a specific group. 

154 

155 Args: 

156 pattern: The group pattern. 

157 key: The proxy key within the group. 

158 """ 

159 group = self._find_group(pattern) (empty)

160 if group and key in group.proxies: (empty)

161 group.healthy_keys.add(key) (empty)

162 

163 def list_routes(self) -> list[dict]: (empty)

164 """List all registered routes with their status. 

165 

166 Returns: 

167 List of dicts with pattern, strategy, proxy_count, healthy_count. 

168 """ 

169 routes = [] (empty)

170 with self._lock: (empty)

171 for pattern, group in self._exact.items(): (empty)

172 routes.append(self._group_info(group)) (empty)

173 for group in self._wildcards: (empty)

174 routes.append(self._group_info(group)) (empty)

175 return routes (empty)

176 

177 def _find_group(self, pattern: str) -> Optional[ProxyGroup]: (empty)

178 """Find a group by its exact pattern string.""" 

179 with self._lock: (empty)

180 if pattern in self._exact: (empty)

181 return self._exact[pattern] (empty)

182 for group in self._wildcards: (empty)

183 if group.pattern == pattern: 183 ↛ 182line 183 didn't jump to line 182 because the condition on line 183 was always true(empty)

184 return group (empty)

185 return None (empty)

186 

187 @staticmethod (empty)

188 def _group_info(group: ProxyGroup) -> dict: (empty)

189 """Build info dict for a proxy group.""" 

190 return { (empty)

191 "pattern": group.pattern, 

192 "strategy": group.strategy.value, 

193 "proxy_count": len(group.proxies), 

194 "healthy_count": len(group.healthy_keys), 

195 "keys": list(group.proxies.keys()), 

196 } 

197 

198 

199def _is_glob_pattern(pattern: str) -> bool: (empty)

200 """Check if a pattern contains glob wildcards.""" 

201 return any(char in pattern for char in ("*", "?", "[")) (empty)