Coverage for src/local_deep_research/web_search_engines/engines/_google_pse_rate_limiter.py: 100%

1"""Process-wide minimum request spacing for Google PSE engine instances.""" 

2 

3import secrets 

4import threading 

5import time 

6 

7_meta_lock = threading.Lock() 

8_scope_locks: dict[str, threading.Lock] = {} 

9_scope_last_request: dict[str, float] = {} 

10_scope_nonce = secrets.token_urlsafe(32) 

11 

12 

13def _scope_key(api_key: str, search_engine_id: str) -> str: 

14 """Return a process-local opaque ID without retaining raw credentials. 

15 

16 Python string hashing is process-keyed SipHash. Two independently ordered 

17 tuple hashes plus a random nonce give the in-memory limiter a compact scope 

18 identity without treating API credentials as password-verification data. 

19 """ 

20 mask = (1 << 64) - 1 

21 first = hash((_scope_nonce, api_key, search_engine_id)) & mask 

22 second = hash((search_engine_id, api_key, _scope_nonce)) & mask 

23 return f"{first:016x}{second:016x}" 

24 

25 

26def _get_scope_lock(scope: str) -> threading.Lock: 

27 with _meta_lock: 

28 lock = _scope_locks.get(scope) 

29 if lock is None: 

30 lock = threading.Lock() 

31 _scope_locks[scope] = lock 

32 return lock 

33 

34 

35def respect_rate_limit( 

36 api_key: str, 

37 search_engine_id: str, 

38 interval_seconds: float, 

39) -> float: 

40 """Sleep as needed to reserve a request slot for one PSE configuration.""" 

41 if interval_seconds <= 0: 

42 return 0.0 

43 

44 scope = _scope_key(api_key, search_engine_id) 

45 lock = _get_scope_lock(scope) 

46 with lock: 

47 elapsed = time.monotonic() - _scope_last_request.get(scope, 0.0) 

48 wait_time = max(0.0, interval_seconds - elapsed) 

49 if wait_time > 0: 

50 time.sleep(wait_time) 

51 _scope_last_request[scope] = time.monotonic() 

52 return wait_time 

53 

54 

55def reset_for_tests() -> None: 

56 """Clear process-wide state. Intended for unit tests only.""" 

57 with _meta_lock: 

58 _scope_locks.clear() 

59 _scope_last_request.clear()