Skip to content

Commit ed682fb

Browse files
committed
feat: add generic tier-aware rate limiting framework
Add RATE_LIMIT_TIERS class attribute and resolve_rate_limiter() static method to BaseProvider. Any provider with subscription tiers can define RATE_LIMIT_TIERS and pass tier + tiers to resolve_rate_limiter() to get automatic tier-aware rate limiter creation. Precedence: tier > explicit rate_limiter > None. Tier matching is case-insensitive. Invalid tiers raise ValueError. This is a provider-agnostic foundation -- no provider-specific code. Providers adopt it by defining RATE_LIMIT_TIERS and calling resolve_rate_limiter() in their constructor. Ref: repowise-dev#68
1 parent adde5e3 commit ed682fb

2 files changed

Lines changed: 127 additions & 1 deletion

File tree

packages/core/src/repowise/core/providers/llm/base.py

Lines changed: 56 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,10 @@
1515

1616
from abc import ABC, abstractmethod
1717
from dataclasses import dataclass, field
18-
from typing import Any, AsyncIterator, Protocol, runtime_checkable
18+
from typing import Any, AsyncIterator, Protocol, TYPE_CHECKING, runtime_checkable
19+
20+
if TYPE_CHECKING:
21+
from repowise.core.rate_limiter import RateLimitConfig, RateLimiter
1922

2023

2124
@dataclass
@@ -59,8 +62,60 @@ class BaseProvider(ABC):
5962
- Return GeneratedResponse with correct token counts
6063
- Raise ProviderError on non-recoverable API errors
6164
- Raise RateLimitError on 429 responses after retries are exhausted
65+
66+
Class Attributes:
67+
RATE_LIMIT_TIERS: Optional mapping of tier name to RateLimitConfig.
68+
Providers with subscription tiers (e.g., Z.AI's lite/pro/max,
69+
MiniMax's starter/plus/max) define this to support tier-aware
70+
rate limiting. When set, users can pass ``tier="pro"`` to the
71+
constructor and the appropriate rate limiter is created automatically.
6272
"""
6373

74+
RATE_LIMIT_TIERS: dict[str, Any] = {} # Override in subclasses
75+
76+
@staticmethod
77+
def resolve_rate_limiter(
78+
tier: str | None = None,
79+
tiers: dict[str, Any] | None = None,
80+
rate_limiter: Any | None = None,
81+
) -> Any | None:
82+
"""Resolve rate limiter using tier precedence.
83+
84+
Precedence: tier > explicit rate_limiter > None.
85+
86+
When tier is set, it takes precedence -- it represents a specific
87+
provider signal that overrides the generic registry default.
88+
89+
Args:
90+
tier: Tier name (e.g., 'lite', 'pro', 'max'). Case-insensitive.
91+
tiers: Mapping of tier name to RateLimitConfig.
92+
rate_limiter: Explicitly provided RateLimiter instance.
93+
94+
Returns:
95+
A RateLimiter instance, or None if neither tier nor
96+
rate_limiter is provided.
97+
98+
Raises:
99+
ValueError: If tier is not found in the tiers mapping.
100+
"""
101+
# Late import to avoid circular dependency at module level
102+
from repowise.core.rate_limiter import RateLimiter
103+
104+
if tier is not None:
105+
if not tiers:
106+
msg = f"Tier {tier!r} specified but provider defines no tiers"
107+
raise ValueError(msg)
108+
tier_key = tier.lower()
109+
tier_config = tiers.get(tier_key)
110+
if tier_config is None:
111+
valid = ", ".join(sorted(tiers))
112+
msg = f"Unknown tier {tier!r}. Valid tiers: {valid}"
113+
raise ValueError(msg)
114+
return RateLimiter(tier_config)
115+
if rate_limiter is not None:
116+
return rate_limiter
117+
return None
118+
64119
@abstractmethod
65120
async def generate(
66121
self,
Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
"""Unit tests for BaseProvider.resolve_rate_limiter (generic tier framework).
2+
3+
These tests verify the tier resolution logic independent of any specific provider.
4+
Any provider that defines RATE_LIMIT_TIERS gets this behavior for free via
5+
BaseProvider.resolve_rate_limiter().
6+
"""
7+
8+
from __future__ import annotations
9+
10+
import pytest
11+
12+
from repowise.core.providers.llm.base import BaseProvider
13+
from repowise.core.rate_limiter import RateLimitConfig, RateLimiter
14+
15+
16+
def test_resolve_rate_limiter_with_tier():
17+
"""resolve_rate_limiter should create a limiter from tier config."""
18+
tiers = {
19+
"basic": RateLimitConfig(requests_per_minute=5, tokens_per_minute=10_000),
20+
"premium": RateLimitConfig(requests_per_minute=50, tokens_per_minute=100_000),
21+
}
22+
limiter = BaseProvider.resolve_rate_limiter(tier="premium", tiers=tiers)
23+
assert limiter is not None
24+
assert limiter.config.requests_per_minute == 50
25+
26+
27+
def test_resolve_rate_limiter_tier_overrides_explicit():
28+
"""Tier should take precedence over explicit rate_limiter."""
29+
tiers = {"pro": RateLimitConfig(requests_per_minute=30, tokens_per_minute=100_000)}
30+
explicit = RateLimiter(RateLimitConfig(requests_per_minute=999, tokens_per_minute=999_999))
31+
limiter = BaseProvider.resolve_rate_limiter(tier="pro", tiers=tiers, rate_limiter=explicit)
32+
assert limiter is not explicit
33+
assert limiter.config.requests_per_minute == 30
34+
35+
36+
def test_resolve_rate_limiter_explicit_without_tier():
37+
"""Without tier, explicit rate_limiter should be returned."""
38+
explicit = RateLimiter(RateLimitConfig(requests_per_minute=42, tokens_per_minute=420_000))
39+
limiter = BaseProvider.resolve_rate_limiter(rate_limiter=explicit)
40+
assert limiter is explicit
41+
42+
43+
def test_resolve_rate_limiter_none_when_nothing_provided():
44+
"""Should return None when neither tier nor rate_limiter is provided."""
45+
limiter = BaseProvider.resolve_rate_limiter()
46+
assert limiter is None
47+
48+
49+
def test_resolve_rate_limiter_invalid_tier():
50+
"""Invalid tier should raise ValueError."""
51+
tiers = {"basic": RateLimitConfig(requests_per_minute=5, tokens_per_minute=10_000)}
52+
with pytest.raises(ValueError, match="Unknown tier"):
53+
BaseProvider.resolve_rate_limiter(tier="enterprise", tiers=tiers)
54+
55+
56+
def test_resolve_rate_limiter_tier_but_no_tiers_defined():
57+
"""Tier with empty tiers dict should raise ValueError."""
58+
with pytest.raises(ValueError, match="defines no tiers"):
59+
BaseProvider.resolve_rate_limiter(tier="pro", tiers={})
60+
61+
62+
def test_resolve_rate_limiter_case_insensitive():
63+
"""Tier matching should be case-insensitive."""
64+
tiers = {"pro": RateLimitConfig(requests_per_minute=30, tokens_per_minute=100_000)}
65+
limiter = BaseProvider.resolve_rate_limiter(tier="PRO", tiers=tiers)
66+
assert limiter.config.requests_per_minute == 30
67+
68+
69+
def test_base_provider_default_empty_tiers():
70+
"""BaseProvider should have empty RATE_LIMIT_TIERS by default."""
71+
assert BaseProvider.RATE_LIMIT_TIERS == {}

0 commit comments

Comments
 (0)