Source code for ai.analysis.model.token_price_calculator

from dataclasses import dataclass

from ai.analysis.model.assistant_chat_model import AssistantChatModel
from ai.analysis.model.token_type import TokenType
from ai.analysis.money.money import Money




[docs]
@dataclass
class ModelPriceInfo:
    """Cost information for input and output tokens of a model.

    Attributes:
        input_token_cost (Money): Cost per input token.
        output_token_cost (Money): Cost per output token.
    """

    input_token_cost: Money
    output_token_cost: Money



[docs]
    def get_cost_for(self, token_type: TokenType) -> Money:
        """Return the cost for the specified token type.

        Args:
            token_type (TokenType): The type of token (INPUT or OUTPUT).

        Returns:
            Money: The cost associated with the token type.

        Raises:
            ValueError: If an unsupported token type is provided.
        """
        if token_type == TokenType.INPUT:
            return self.input_token_cost
        elif token_type == TokenType.OUTPUT:
            return self.output_token_cost
        raise ValueError("Wrong token type")








[docs]
class TokenPriceCalculator:
    """Calculator for determining token costs of supported assistant models."""

    model_price_infos: dict[AssistantChatModel, ModelPriceInfo] = {
        AssistantChatModel.GPT_4O: ModelPriceInfo(Money(2.5e-6), Money(10e-6)),
        AssistantChatModel.GPT_4O_MINI: ModelPriceInfo(Money(0.15e-6), Money(0.6e-6)),
        AssistantChatModel.GPT_4_1: ModelPriceInfo(Money(2e-6), Money(8e-6)),
        AssistantChatModel.GPT_4_1_MINI: ModelPriceInfo(Money(0.40e-6), Money(1.60e-6)),
        AssistantChatModel.GPT_4_1_NANO: ModelPriceInfo(Money(0.1e-6), Money(0.4e-6)),
        AssistantChatModel.LLAMA_70B: ModelPriceInfo(Money(0.12e-6), Money(0.3e-6)),
        AssistantChatModel.QWEN_2_5_7B: ModelPriceInfo(Money(0.025e-6), Money(0.05e-6)),
        AssistantChatModel.LLAMA_3B: ModelPriceInfo(Money(0.015e-6), Money(0.025e-6)),
        AssistantChatModel.QWEN_2_5_7B_TOGETHER: ModelPriceInfo(Money(0.3e-6), Money(0.3e-6)),
        AssistantChatModel.DEEP_SEEK_R1: ModelPriceInfo(Money(0.5e-6), Money(2.15e-6)),
    }



[docs]
    def get_cost(self, model: AssistantChatModel, token_type: TokenType) -> Money:
        """Calculate the cost for a given model and token type.

        Args:
            model (AssistantChatModel): The chat model to query.
            token_type (TokenType): The type of token (INPUT or OUTPUT).

        Returns:
            Money: The cost for one token of the specified type for the given model.

        Raises:
            ValueError: If the model is not found in the price information.
        """
        if model in self.model_price_infos:
            return self.model_price_infos[model].get_cost_for(token_type)
        raise ValueError(f"Model not found in Token Price Calculator {model.value}")