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
« 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"""
5from dataclasses import dataclass, field
6from enum import Enum
8from pydantic import BaseModel, Field
11class Provider(Enum):
12 """Supported API providers."""
14 OPENAI = "openai"
15 ANTHROPIC = "anthropic"
16 OPENROUTER = "openrouter"
19@dataclass
20class RateLimitInfo:
21 """Information about current rate limits."""
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
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
33 # Retry information
34 retry_after: float | None = None # Seconds to wait
36 # Provider-specific metadata
37 metadata: dict[str, str] = field(default_factory=dict)
40class ProviderConfig(BaseModel):
41 """Configuration for API provider rate limit handling."""
43 provider: Provider
44 base_url: str
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
51 token_limit_header: str | None = None
52 token_remaining_header: str | None = None
53 token_reset_header: str | None = None
55 retry_after_header: str | None = None
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
61 # Rate limit discovery
62 supports_dynamic_limits: bool = True
63 auth_endpoint: str | None = None # For checking limits via API
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)
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)
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}
107def get_provider_config(provider: Provider) -> ProviderConfig:
108 """Get configuration for a specific provider."""
109 return PROVIDER_CONFIGS[provider]
112def detect_provider_from_url(url: str) -> Provider | None:
113 """Detect provider from API URL."""
114 url_lower = url.lower()
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
123 return None
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()
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))
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))
150 # Extract retry information
151 if config.retry_after_header:
152 info.retry_after = _safe_float(headers.get(config.retry_after_header))
154 # Store all headers as metadata
155 info.metadata = dict(headers)
157 return info
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
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