neollm/myllm/abstract_myllm.py

from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Generator, Generic, Optional, TypeAlias, cast

from neollm.myllm.print_utils import print_inputs, print_metadata, print_outputs
from neollm.types import (
    InputType,
    OutputType,
    PriceInfo,
    StreamOutputType,
    TimeInfo,
    TokenInfo,
)
from neollm.utils.utils import cprint

if TYPE_CHECKING:
    from typing import Any

    from neollm.myllm.myl3m2 import MyL3M2

    _MyL3M2: TypeAlias = MyL3M2[Any, Any]


class AbstractMyLLM(ABC, Generic[InputType, OutputType]):
    """MyLLM, MyL3M2の抽象クラス"""

    inputs: InputType | None
    outputs: OutputType | None
    silent_set: set[str]
    verbose: bool
    time: float = 0.0
    time_detail: TimeInfo = TimeInfo()
    parent: Optional["_MyL3M2"] = None
    do_stream: bool

    @property
    @abstractmethod
    def token(self) -> TokenInfo:
        """LLMの利用トークン数

        Returns:
            TokenInfo: トークン数 (入力, 出力, 合計)
            >>> TokenInfo(input=1588, output=128, total=1716)
        """

    @property
    def custom_token(self) -> TokenInfo | None:
        """料金計算用トークン（Gemini用）"""
        return None

    @property
    @abstractmethod
    def price(self) -> PriceInfo:
        """LLMの利用料金 (USD)

        Returns:
            PriceInfo: 利用料金 (USD) (入力, 出力, 合計)
            >>> PriceInfo(input=0.002382, output=0.000256, total=0.002638)
        """

    @abstractmethod
    def _call(self, inputs: InputType, stream: bool = False) -> Generator[StreamOutputType, None, OutputType]:
        """MyLLMの子クラスのメインロジック

        streamとnon-streamの両方のコードを書く必要がある

        Args:
            inputs (InputType): LLMへの入力
            stream (bool, optional): streamの有無. Defaults to False.

        Yields:
            Generator[StreamOutputType, None, OutputType]: LLMのstream出力

        Returns:
            OutputType: LLMの出力
        """

    def __call__(self, inputs: InputType) -> OutputType:
        """MyLLMのメインロジック

        Args:
            inputs (InputType): LLMへの入力

        Returns:
            OutputType: LLMの出力
        """
        it: Generator[StreamOutputType, None, OutputType] = self._call(inputs, stream=self.do_stream)
        while True:
            try:
                next(it)
            except StopIteration as e:
                outputs = cast(OutputType, e.value)
                return outputs
            except Exception as e:
                raise e

    def call_stream(self, inputs: InputType) -> Generator[StreamOutputType, None, OutputType]:
        """MyLLMのメインロジック(stream処理)

        Args:
            inputs (InputType): LLMへの入力

        Yields:
            Generator[StreamOutputType, None, OutputType]: LLMのstream出力

        Returns:
            LLMの出力
        """
        it: Generator[StreamOutputType, None, OutputType] = self._call(inputs, stream=True)
        while True:
            try:
                delta_content = next(it)
                yield delta_content
            except StopIteration as e:
                outputs = cast(OutputType, e.value)
                return outputs
            except Exception as e:
                raise e

    def _print_inputs(self) -> None:
        if self.inputs is None:
            return
        if not ("inputs" not in self.silent_set and self.verbose):
            return
        print_inputs(self.inputs)

    def _print_outputs(self) -> None:
        if self.outputs is None:
            return
        if not ("outputs" not in self.silent_set and self.verbose):
            return
        print_outputs(self.outputs)

    def _print_metadata(self) -> None:
        if not ("metadata" not in self.silent_set and self.verbose):
            return
        print_metadata(self.time, self.token, self.price)

    def _print_start(self, sep: str = "-") -> None:
        if not self.verbose:
            return
        if self.parent is None:
            cprint("PARENT", color="red", background=True)
        print(self, sep * (99 - len(str(self))))

    def _print_end(self, sep: str = "-") -> None:
        if not self.verbose:
            return
        print(sep * 100)