src/vsp/llm/openai/openai_model.py

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)