Coverage for src/chat_limiter/providers.py: 100%

80 statements  

« prev     ^ index     » next       coverage.py v7.9.2, created at 2025-07-11 14:11 +0100

1""" 

2Provider-specific configurations and rate limit header mappings. 

3""" 

4 

5from dataclasses import dataclass, field 

6from enum import Enum 

7 

8from pydantic import BaseModel, Field 

9 

10 

11class Provider(Enum): 

12 """Supported API providers.""" 

13 

14 OPENAI = "openai" 

15 ANTHROPIC = "anthropic" 

16 OPENROUTER = "openrouter" 

17 

18 

19@dataclass 

20class RateLimitInfo: 

21 """Information about current rate limits.""" 

22 

23 # Request limits 

24 requests_limit: int | None = None 

25 requests_remaining: int | None = None 

26 requests_reset: int | float | None = None # Unix timestamp or seconds 

27 

28 # Token limits 

29 tokens_limit: int | None = None 

30 tokens_remaining: int | None = None 

31 tokens_reset: int | float | None = None # Unix timestamp or seconds 

32 

33 # Retry information 

34 retry_after: float | None = None # Seconds to wait 

35 

36 # Provider-specific metadata 

37 metadata: dict[str, str] = field(default_factory=dict) 

38 

39 

40class ProviderConfig(BaseModel): 

41 """Configuration for API provider rate limit handling.""" 

42 

43 provider: Provider 

44 base_url: str 

45 

46 # Rate limit header mappings 

47 request_limit_header: str | None = None 

48 request_remaining_header: str | None = None 

49 request_reset_header: str | None = None 

50 

51 token_limit_header: str | None = None 

52 token_remaining_header: str | None = None 

53 token_reset_header: str | None = None 

54 

55 retry_after_header: str | None = None 

56 

57 # Rate limits (must be provided by user or discovered from API) 

58 default_request_limit: int | None = None 

59 default_token_limit: int | None = None 

60 

61 # Rate limit discovery 

62 supports_dynamic_limits: bool = True 

63 auth_endpoint: str | None = None # For checking limits via API 

64 

65 # Retry configuration (no defaults - must be specified by user if needed) 

66 max_retries: int | None = None 

67 base_backoff: float | None = None 

68 max_backoff: float = Field(default=60.0, ge=1.0) 

69 

70 # Safety buffers 

71 request_buffer_ratio: float = Field(default=0.9, ge=0.1, le=1.0) 

72 token_buffer_ratio: float = Field(default=0.9, ge=0.1, le=1.0) 

73 

74 

75# Provider-specific configurations 

76PROVIDER_CONFIGS = { 

77 Provider.OPENAI: ProviderConfig( 

78 provider=Provider.OPENAI, 

79 base_url="https://api.openai.com/v1", 

80 request_limit_header="x-ratelimit-limit-requests", 

81 request_remaining_header="x-ratelimit-remaining-requests", 

82 request_reset_header="x-ratelimit-reset-requests", 

83 token_limit_header="x-ratelimit-limit-tokens", 

84 token_remaining_header="x-ratelimit-remaining-tokens", 

85 token_reset_header="x-ratelimit-reset-tokens", 

86 retry_after_header="retry-after", 

87 supports_dynamic_limits=True, 

88 ), 

89 Provider.ANTHROPIC: ProviderConfig( 

90 provider=Provider.ANTHROPIC, 

91 base_url="https://api.anthropic.com/v1", 

92 request_remaining_header="anthropic-ratelimit-requests-remaining", 

93 token_limit_header="anthropic-ratelimit-tokens-limit", 

94 token_reset_header="anthropic-ratelimit-tokens-reset", 

95 retry_after_header="retry-after", 

96 supports_dynamic_limits=True, 

97 ), 

98 Provider.OPENROUTER: ProviderConfig( 

99 provider=Provider.OPENROUTER, 

100 base_url="https://openrouter.ai/api/v1", 

101 auth_endpoint="https://openrouter.ai/api/v1/auth/key", 

102 supports_dynamic_limits=True, 

103 ), 

104} 

105 

106 

107def get_provider_config(provider: Provider) -> ProviderConfig: 

108 """Get configuration for a specific provider.""" 

109 return PROVIDER_CONFIGS[provider] 

110 

111 

112def detect_provider_from_url(url: str) -> Provider | None: 

113 """Detect provider from API URL.""" 

114 url_lower = url.lower() 

115 

116 if "openai.com" in url_lower: 

117 return Provider.OPENAI 

118 elif "anthropic.com" in url_lower: 

119 return Provider.ANTHROPIC 

120 elif "openrouter.ai" in url_lower: 

121 return Provider.OPENROUTER 

122 

123 return None 

124 

125 

126def extract_rate_limit_info( 

127 headers: dict[str, str], config: ProviderConfig 

128) -> RateLimitInfo: 

129 """Extract rate limit information from response headers.""" 

130 info = RateLimitInfo() 

131 

132 # Extract request limits 

133 if config.request_limit_header: 

134 info.requests_limit = _safe_int(headers.get(config.request_limit_header)) 

135 if config.request_remaining_header: 

136 info.requests_remaining = _safe_int( 

137 headers.get(config.request_remaining_header) 

138 ) 

139 if config.request_reset_header: 

140 info.requests_reset = _safe_float(headers.get(config.request_reset_header)) 

141 

142 # Extract token limits 

143 if config.token_limit_header: 

144 info.tokens_limit = _safe_int(headers.get(config.token_limit_header)) 

145 if config.token_remaining_header: 

146 info.tokens_remaining = _safe_int(headers.get(config.token_remaining_header)) 

147 if config.token_reset_header: 

148 info.tokens_reset = _safe_float(headers.get(config.token_reset_header)) 

149 

150 # Extract retry information 

151 if config.retry_after_header: 

152 info.retry_after = _safe_float(headers.get(config.retry_after_header)) 

153 

154 # Store all headers as metadata 

155 info.metadata = dict(headers) 

156 

157 return info 

158 

159 

160def _safe_int(value: str | None) -> int | None: 

161 """Safely convert string to int.""" 

162 if value is None: 

163 return None 

164 try: 

165 return int(value) 

166 except (ValueError, TypeError): 

167 return None 

168 

169 

170def _safe_float(value: str | None) -> float | None: 

171 """Safely convert string to float.""" 

172 if value is None: 

173 return None 

174 try: 

175 return float(value) 

176 except (ValueError, TypeError): 

177 return None