Coverage for src / local_deep_research / database / models / rate_limiting.py: 93%

1""" 

2Rate limiting tracking models. 

3""" 

4 

5from sqlalchemy import Boolean, Column, Float, Index, Integer, String 

6from sqlalchemy_utc import UtcDateTime, utcnow 

7 

8from .base import Base 

9 

10 

11class RateLimitAttempt(Base): 

12 """ 

13 Track individual rate limit retry attempts. 

14 """ 

15 

16 __tablename__ = "rate_limit_attempts" 

17 

18 id = Column(Integer, primary_key=True, index=True) 

19 engine_type = Column(String(100), nullable=False, index=True) 

20 timestamp = Column(Float, nullable=False, index=True) 

21 wait_time = Column(Float, nullable=False) 

22 retry_count = Column(Integer, nullable=False) 

23 success = Column(Boolean, nullable=False) 

24 error_type = Column(String(100), nullable=True) 

25 created_at = Column(UtcDateTime, server_default=utcnow(), nullable=False) 

26 

27 # Compound indexes for query performance optimization 

28 __table_args__ = ( 

29 Index("idx_ratelimit_engine_timestamp", "engine_type", "timestamp"), 

30 Index("idx_ratelimit_success_timestamp", "success", "timestamp"), 

31 Index( 

32 "idx_ratelimit_engine_success", 

33 "engine_type", 

34 "success", 

35 "timestamp", 

36 ), 

37 ) 

38 

39 def __repr__(self): 

40 return f"<RateLimitAttempt(engine={self.engine_type}, success={self.success}, wait={self.wait_time}s)>" 

41 

42 

43class RateLimitEstimate(Base): 

44 """ 

45 Store current rate limit estimates per engine. 

46 """ 

47 

48 __tablename__ = "rate_limit_estimates" 

49 

50 id = Column(Integer, primary_key=True, index=True) 

51 engine_type = Column(String(100), nullable=False, unique=True, index=True) 

52 base_wait_seconds = Column(Float, nullable=False) 

53 min_wait_seconds = Column(Float, nullable=False) 

54 max_wait_seconds = Column(Float, nullable=False) 

55 last_updated = Column(Float, nullable=False) 

56 total_attempts = Column(Integer, default=0, nullable=False) 

57 success_rate = Column(Float, default=0.0, nullable=False) 

58 created_at = Column(UtcDateTime, server_default=utcnow(), nullable=False) 

59 updated_at = Column( 

60 UtcDateTime, server_default=utcnow(), onupdate=utcnow(), nullable=False 

61 ) 

62 

63 def __repr__(self): 

64 return f"<RateLimitEstimate(engine={self.engine_type}, wait={self.base_wait_seconds}s, success_rate={self.success_rate:.1%})>"