File size: 2,036 Bytes
3b993c4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 |
from dataclasses import dataclass
from enum import Enum
@dataclass
class RateLimit:
"""
Dataclass representing rate limit information for OpenAI models.
Attributes:
requests_per_minute (int): The maximum number of requests allowed per minute.
tokens_per_minute (int): The maximum number of tokens allowed per minute.
"""
requests_per_minute: int
tokens_per_minute: int
class OpenAIModel(Enum):
"""
Enum representing different OpenAI models with their rate limits.
Attributes:
GPT_4_O: Represents the GPT-4 model.
GPT_4_MINI: Represents the GPT-4 mini model.
Methods:
get_rate_limit(model: OpenAIModel) -> RateLimit:
Get the rate limit for a specific model.
"""
GPT_4_O = "gpt-4o"
GPT_4_MINI = "gpt-4o-mini"
@classmethod
def get_rate_limit(cls, model: "OpenAIModel") -> RateLimit:
"""
Get the rate limit for a specific OpenAI model.
Args:
model (OpenAIModel): The OpenAI model to get the rate limit for.
Returns:
RateLimit: The rate limit information for the specified model.
"""
return cls._rate_limits()[model]
@classmethod
def _rate_limits(cls) -> dict[Enum, RateLimit]:
"""
Define the rate limits for each OpenAI model.
Returns:
dict[Enum, RateLimit]: A dictionary mapping each model to its rate limit.
"""
return {
cls.GPT_4_O: RateLimit(500, 30_000),
cls.GPT_4_MINI: RateLimit(500, 60_000),
}
def get_openai_model_rate_limit(model: OpenAIModel) -> RateLimit:
"""
Get the rate limit for a specific OpenAI model.
This is a convenience function that calls the get_rate_limit method of the OpenAIModel enum.
Args:
model (OpenAIModel): The OpenAI model to get the rate limit for.
Returns:
RateLimit: The rate limit information for the specified model.
"""
return OpenAIModel.get_rate_limit(model)
|