Coverage for src/local_deep_research/web_search_engines/rate_limiting/llm/wrapper.py: 95%

1""" 

2Rate-limited wrapper for LLM calls. 

3""" 

4 

5from typing import Optional 

6from urllib.parse import urlparse 

7 

8from ....security import scrub_error 

9from ....security.secure_logging import logger 

10from tenacity import ( 

11 retry, 

12 stop_after_attempt, 

13 retry_if_exception, 

14) 

15from tenacity.wait import wait_base 

16 

17from ....llm.providers.base import normalize_provider 

18 

19from ..tracker import get_tracker 

20from ..exceptions import RateLimitError 

21from .detection import is_llm_rate_limit_error, extract_retry_after 

22 

23 

24class AdaptiveLLMWait(wait_base): 

25 """Adaptive wait strategy for LLM rate limiting.""" 

26 

27 def __init__(self, tracker, engine_type: str): 

28 self.tracker = tracker 

29 self.engine_type = engine_type 

30 self.last_error = None 

31 

32 def __call__(self, retry_state) -> float: 

33 # Store the error for potential retry-after extraction 

34 if retry_state.outcome and retry_state.outcome.failed: 

35 self.last_error = retry_state.outcome.exception() 

36 

37 # Get adaptive wait time from tracker 

38 wait_time: float = self.tracker.get_wait_time(self.engine_type) 

39 

40 # If we have a retry-after from the error, use it 

41 if self.last_error: 

42 retry_after = extract_retry_after(self.last_error) 

43 if retry_after > 0: 

44 wait_time = max(wait_time, float(retry_after)) 

45 

46 logger.info( 

47 f"LLM rate limit wait for {self.engine_type}: {wait_time:.2f}s" 

48 ) 

49 return wait_time 

50 

51 

52def create_rate_limited_llm_wrapper(base_llm, provider: Optional[str] = None): 

53 """ 

54 Create a rate-limited wrapper around an LLM instance. 

55 

56 Args: 

57 base_llm: The base LLM instance to wrap 

58 provider: Optional provider name (e.g., 'openai', 'anthropic') 

59 

60 Returns: 

61 A wrapped LLM instance with rate limiting capabilities 

62 """ 

63 

64 class RateLimitedLLMWrapper: 

65 """Wrapper that adds rate limiting to LLM calls.""" 

66 

67 def __init__(self, llm, provider_name: Optional[str] = None): 

68 self.base_llm = llm 

69 self.provider = provider_name 

70 self.rate_limiter = None 

71 

72 # Only setup rate limiting if enabled 

73 if self._should_rate_limit(): 73 ↛ 74line 73 didn't jump to line 74 because the condition on line 73 was never true

74 self.rate_limiter = get_tracker() 

75 logger.info( 

76 f"Rate limiting enabled for LLM provider: {self._get_rate_limit_key()}" 

77 ) 

78 

79 def _should_rate_limit(self) -> bool: 

80 """Check if rate limiting should be applied to this LLM.""" 

81 # Rate limiting for LLMs is currently disabled by default 

82 # TODO: Pass settings_snapshot to enable proper configuration 

83 return False 

84 

85 def _check_if_local_model(self) -> bool: 

86 """Check if the LLM is a local model that shouldn't be rate limited.""" 

87 # Don't rate limit local models 

88 local_providers = [ 

89 "ollama", 

90 "lmstudio", 

91 "llamacpp", 

92 "local", 

93 "none", 

94 ] 

95 if normalize_provider(self.provider) in local_providers: 

96 logger.debug( 

97 f"Skipping rate limiting for local provider: {self.provider}" 

98 ) 

99 return True 

100 

101 # Check if base URL indicates local model 

102 if hasattr(self.base_llm, "base_url"): 

103 base_url = str(self.base_llm.base_url) 

104 if any( 

105 local in base_url 

106 for local in ["localhost", "127.0.0.1", "0.0.0.0"] 

107 ): 

108 logger.debug( 

109 f"Skipping rate limiting for local URL: {base_url}" 

110 ) 

111 return True 

112 

113 return False 

114 

115 def _get_rate_limit_key(self) -> str: 

116 """Build composite key: provider-url-model""" 

117 provider = self.provider or "unknown" 

118 

119 # Extract URL 

120 url = "unknown" 

121 if hasattr(self.base_llm, "base_url"): 

122 url = str(self.base_llm.base_url) 

123 elif hasattr(self.base_llm, "_client") and hasattr( 

124 self.base_llm._client, "base_url" 

125 ): 

126 url = str(self.base_llm._client.base_url) 

127 

128 # Clean URL: remove protocol and trailing slashes 

129 if url != "unknown": 

130 parsed = urlparse(url) 

131 url = parsed.netloc or parsed.path 

132 url = url.rstrip("/") 

133 

134 # Extract model 

135 model = "unknown" 

136 if hasattr(self.base_llm, "model_name"): 

137 model = str(self.base_llm.model_name) 

138 elif hasattr(self.base_llm, "model"): 

139 model = str(self.base_llm.model) 

140 

141 # Clean model name 

142 model = model.replace("/", "-").replace(":", "-") 

143 

144 return f"{provider}-{url}-{model}" 

145 

146 def invoke(self, *args, **kwargs): 

147 """Invoke the LLM with rate limiting if enabled.""" 

148 if self.rate_limiter: 

149 rate_limit_key = self._get_rate_limit_key() 

150 

151 # Define retry logic 

152 @retry( 

153 wait=AdaptiveLLMWait(self.rate_limiter, rate_limit_key), 

154 stop=stop_after_attempt(3), 

155 retry=retry_if_exception(is_llm_rate_limit_error), 

156 ) 

157 def _invoke_with_retry(): 

158 return self._do_invoke(*args, **kwargs) 

159 

160 try: 

161 result = _invoke_with_retry() 

162 

163 # Record successful attempt 

164 self.rate_limiter.record_outcome( 

165 engine_type=rate_limit_key, 

166 wait_time=0, # First attempt had no wait 

167 success=True, 

168 retry_count=0, 

169 ) 

170 

171 return result 

172 

173 except Exception as e: 

174 # Only record rate limit failures, not general failures 

175 if is_llm_rate_limit_error(e): 175 ↛ 182line 175 didn't jump to line 182 because the condition on line 175 was always true

176 self.rate_limiter.record_outcome( 

177 engine_type=rate_limit_key, 

178 wait_time=0, 

179 success=False, 

180 retry_count=0, 

181 ) 

182 raise 

183 else: 

184 # No rate limiting, just invoke directly 

185 return self._do_invoke(*args, **kwargs) 

186 

187 def _do_invoke(self, *args, **kwargs): 

188 """Actually invoke the LLM.""" 

189 try: 

190 return self.base_llm.invoke(*args, **kwargs) 

191 except Exception as e: 

192 # Check if it's a rate limit error and wrap it 

193 if is_llm_rate_limit_error(e): 

194 logger.warning("LLM rate limit error detected") 

195 # Scrub before wrapping: LLM providers can echo the 

196 # Authorization header back in a 429 body. The scrub 

197 # keeps non-credential text (e.g. "retry after 20 

198 # seconds") intact for extract_retry_after(). `from 

199 # None` suppresses the __context__ chain, which still 

200 # carries the raw message. 

201 safe_msg = scrub_error(e) 

202 raise RateLimitError( 

203 f"LLM rate limit: {safe_msg}" 

204 ) from None 

205 raise 

206 

207 # Pass through any other attributes to the base LLM 

208 def __getattr__(self, name): 

209 return getattr(self.base_llm, name) 

210 

211 def close(self): 

212 """Close underlying HTTP clients held by this LLM. Idempotent.""" 

213 try: 

214 from ....utilities.llm_utils import _close_base_llm 

215 

216 _close_base_llm(self.base_llm) 

217 except Exception: 

218 logger.debug( 

219 "best-effort cleanup of HTTP clients on shutdown", 

220 exc_info=True, 

221 ) 

222 

223 def __str__(self): 

224 return f"RateLimited({str(self.base_llm)})" 

225 

226 def __repr__(self): 

227 return f"RateLimitedLLMWrapper({repr(self.base_llm)})" 

228 

229 return RateLimitedLLMWrapper(base_llm, provider)