neollm/types/info.py

from typing import Literal

from neollm.types._model import DictableBaseModel

PrintColor = Literal["black", "red", "green", "yellow", "blue", "magenta", "cyan", "white"]


class TimeInfo(DictableBaseModel):
    total: float = 0.0
    """処理時間合計 preprocess + main + postprocess"""
    preprocess: float = 0.0
    """前処理時間"""
    main: float = 0.0
    """メイン処理時間"""
    postprocess: float = 0.0
    """後処理時間"""

    def __repr__(self) -> str:
        return (
            f"TimeInfo(total={self.total:.3f}, preprocess={self.preprocess:.3f}, main={self.main:.3f}, "
            f"postprocess={self.postprocess:.3f})"
        )


class TokenInfo(DictableBaseModel):
    input: int
    """入力部分のトークン数"""
    output: int
    """出力部分のトークン数"""
    total: int
    """合計トークン数"""

    def __add__(self, other: "TokenInfo") -> "TokenInfo":
        if not isinstance(other, TokenInfo):
            raise TypeError(f"{other} is not TokenInfo")
        return TokenInfo(
            input=self.input + other.input, output=self.output + other.output, total=self.total + other.total
        )

    def __iadd__(self, other: "TokenInfo") -> "TokenInfo":
        if not isinstance(other, TokenInfo):
            raise TypeError(f"{other} is not TokenInfo")
        self.input += other.input
        self.output += other.output
        self.total += other.total
        return self


class PriceInfo(DictableBaseModel):
    input: float
    """入力部分の費用 (USD)"""
    output: float
    """出力部分の費用 (USD)"""
    total: float
    """合計費用 (USD)"""

    def __add__(self, other: "PriceInfo") -> "PriceInfo":
        if not isinstance(other, PriceInfo):
            raise TypeError(f"{other} is not PriceInfo")
        return PriceInfo(
            input=self.input + other.input, output=self.output + other.output, total=self.total + other.total
        )

    def __iadd__(self, other: "PriceInfo") -> "PriceInfo":
        if not isinstance(other, PriceInfo):
            raise TypeError(f"{other} is not PriceInfo")
        self.input += other.input
        self.output += other.output
        self.total += other.total
        return self

    def __repr__(self) -> str:
        return f"PriceInfo(input={self.input:.3f}, output={self.output:.3f}, total={self.total:.3f})"


class APIPricing(DictableBaseModel):
    """APIの価格設定に関する情報を表すクラス。"""

    input: float
    """入力 1k tokens 当たりのAPI利用料 (USD)"""
    output: float
    """出力 1k tokens 当たりのAPI利用料 (USD)"""