tokenizers

@define()
class AmazonBedrockTokenizer(BaseTokenizer):
    MODEL_PREFIXES_TO_MAX_INPUT_TOKENS = {
        "anthropic.claude-3": 200000,
        "anthropic.claude-v2:1": 200000,
        "anthropic.claude": 100000,
        "cohere.command-r": 128000,
        "cohere.embed": 512,
        "cohere.command": 4000,
        "cohere": 1024,
        "ai21": 8192,
        "meta.llama3-8b-instruct": 8000,
        "meta.llama3-70b-instruct": 8000,
        "meta.llama3-2-1b-instruct": 131000,
        "meta.llama3-2-3b-instruct": 131000,
        "meta.llama3": 128000,
        "mistral.large-2407": 128000,
        "mistral.mistral": 32000,
        "mistral.mixtral": 32000,
        "amazon.nova-micro-v1": 128000,
        "amazon.nova": 300000,
        "amazon.titan-embed-image": 128000,
        "amazon.titan-embed-text": 8000,
        "amazon.titan-text-express-v1": 8000,
        "amazon.titan-text-lite-v1": 4000,
        "amazon.titan-text-premier-v1": 32000,
    }
    MODEL_PREFIXES_TO_MAX_OUTPUT_TOKENS = {
        "anthropic.claude-3-7": 8192,
        "anthropic.claude-3-5": 8192,
        "anthropic.claude": 4096,
        "cohere": 4096,
        "ai21.j2": 8191,
        "meta": 2048,
        "amazon.titan-text-lite": 4096,
        "amazon.titan-text-express": 8192,
        "amazon.titan-text-premier": 3072,
        "amazon.nova": 5000,
        "mistral.mistral": 8192,
        "mistral.mixtral": 4096,
    }

    model: str = field(kw_only=True)
    characters_per_token: int = field(default=4, kw_only=True)

    def count_tokens(self, text: str) -> int:
        return (len(text) + self.characters_per_token - 1) // self.characters_per_token

def count_tokens(self, text: str) -> int:
    return (len(text) + self.characters_per_token - 1) // self.characters_per_token

@define()
class AnthropicTokenizer(BaseTokenizer):
    MODEL_PREFIXES_TO_MAX_INPUT_TOKENS = {"claude-3": 200000, "claude-2.1": 200000, "claude": 100000}
    MODEL_PREFIXES_TO_MAX_OUTPUT_TOKENS = {"claude": 4096}

    client: Anthropic = field(
        default=Factory(lambda: import_optional_dependency("anthropic").Anthropic()),
        kw_only=True,
    )

    def count_tokens(self, text: str | list[BetaMessageParam]) -> int:
        types = import_optional_dependency("anthropic.types.beta")

        # TODO: Refactor all Tokenizers to support Prompt Stack as an input.
        messages = [types.BetaMessageParam(role="user", content=text)] if isinstance(text, str) else text

        usage = self.client.beta.messages.count_tokens(
            model=self.model,
            messages=messages,
        )

        return usage.input_tokens

def count_tokens(self, text: str | list[BetaMessageParam]) -> int:
    types = import_optional_dependency("anthropic.types.beta")

    # TODO: Refactor all Tokenizers to support Prompt Stack as an input.
    messages = [types.BetaMessageParam(role="user", content=text)] if isinstance(text, str) else text

    usage = self.client.beta.messages.count_tokens(
        model=self.model,
        messages=messages,
    )

    return usage.input_tokens

@define()
class BaseTokenizer(ABC, SerializableMixin):
    DEFAULT_MAX_INPUT_TOKENS = 4096
    DEFAULT_MAX_OUTPUT_TOKENS = 1000
    MODEL_PREFIXES_TO_MAX_INPUT_TOKENS = {}
    MODEL_PREFIXES_TO_MAX_OUTPUT_TOKENS = {}

    model: str = field(kw_only=True, metadata={"serializable": True})
    stop_sequences: list[str] = field(default=Factory(list), kw_only=True, metadata={"serializable": True})
    _max_input_tokens: Optional[int] = field(
        kw_only=True, default=None, alias="max_input_tokens", metadata={"serializable": True}
    )
    _max_output_tokens: Optional[int] = field(
        kw_only=True, default=None, alias="max_output_tokens", metadata={"serializable": True}
    )

    @lazy_property()
    def max_input_tokens(self) -> int:
        return self._default_max_input_tokens()

    @lazy_property()
    def max_output_tokens(self) -> int:
        return self._default_max_output_tokens()

    def __attrs_post_init__(self) -> None:
        if self.model is not None:
            if self.max_input_tokens is None:
                self.max_input_tokens = self._default_max_input_tokens()

            if self.max_output_tokens is None:
                # TODO(collin): https://github.com/griptape-ai/griptape/issues/1844
                self.max_output_tokens = self._default_max_output_tokens()

    def count_input_tokens_left(self, text: str) -> int:
        diff = self.max_input_tokens - self.count_tokens(text)

        if diff > 0:
            return diff
        return 0

    def count_output_tokens_left(self, text: str) -> int:
        diff = self.max_output_tokens - self.count_tokens(text)

        if diff > 0:
            return diff
        return 0

    @abstractmethod
    def count_tokens(self, text: str) -> int: ...

    def _default_max_input_tokens(self) -> int:
        tokens = next(
            (
                max_tokens
                for model_prefix, max_tokens in self.MODEL_PREFIXES_TO_MAX_INPUT_TOKENS.items()
                if model_prefix in self.model
            ),
            None,
        )

        if tokens is None:
            logging.warning(
                "Model %s not found in MODEL_PREFIXES_TO_MAX_INPUT_TOKENS, using default value of %s.",
                self.model,
                self.DEFAULT_MAX_INPUT_TOKENS,
            )
            return self.DEFAULT_MAX_INPUT_TOKENS
        return tokens

    def _default_max_output_tokens(self) -> int:
        tokens = next(
            (
                max_tokens
                for model_prefix, max_tokens in self.MODEL_PREFIXES_TO_MAX_OUTPUT_TOKENS.items()
                if model_prefix in self.model
            ),
            None,
        )

        if tokens is None:
            logging.debug(
                "Model %s not found in MODEL_PREFIXES_TO_MAX_OUTPUT_TOKENS, using default value of %s.",
                self.model,
                self.DEFAULT_MAX_OUTPUT_TOKENS,
            )
            return self.DEFAULT_MAX_OUTPUT_TOKENS
        return tokens

def __attrs_post_init__(self) -> None:
    if self.model is not None:
        if self.max_input_tokens is None:
            self.max_input_tokens = self._default_max_input_tokens()

        if self.max_output_tokens is None:
            # TODO(collin): https://github.com/griptape-ai/griptape/issues/1844
            self.max_output_tokens = self._default_max_output_tokens()

def _default_max_input_tokens(self) -> int:
    tokens = next(
        (
            max_tokens
            for model_prefix, max_tokens in self.MODEL_PREFIXES_TO_MAX_INPUT_TOKENS.items()
            if model_prefix in self.model
        ),
        None,
    )

    if tokens is None:
        logging.warning(
            "Model %s not found in MODEL_PREFIXES_TO_MAX_INPUT_TOKENS, using default value of %s.",
            self.model,
            self.DEFAULT_MAX_INPUT_TOKENS,
        )
        return self.DEFAULT_MAX_INPUT_TOKENS
    return tokens

def _default_max_output_tokens(self) -> int:
    tokens = next(
        (
            max_tokens
            for model_prefix, max_tokens in self.MODEL_PREFIXES_TO_MAX_OUTPUT_TOKENS.items()
            if model_prefix in self.model
        ),
        None,
    )

    if tokens is None:
        logging.debug(
            "Model %s not found in MODEL_PREFIXES_TO_MAX_OUTPUT_TOKENS, using default value of %s.",
            self.model,
            self.DEFAULT_MAX_OUTPUT_TOKENS,
        )
        return self.DEFAULT_MAX_OUTPUT_TOKENS
    return tokens

def count_input_tokens_left(self, text: str) -> int:
    diff = self.max_input_tokens - self.count_tokens(text)

    if diff > 0:
        return diff
    return 0

def count_output_tokens_left(self, text: str) -> int:
    diff = self.max_output_tokens - self.count_tokens(text)

    if diff > 0:
        return diff
    return 0

@abstractmethod
def count_tokens(self, text: str) -> int: ...

@lazy_property()
def max_input_tokens(self) -> int:
    return self._default_max_input_tokens()

@lazy_property()
def max_output_tokens(self) -> int:
    return self._default_max_output_tokens()

@define()
class CohereTokenizer(BaseTokenizer):
    MODEL_PREFIXES_TO_MAX_INPUT_TOKENS = {"command-r": 128000, "command": 4096, "embed": 512}
    MODEL_PREFIXES_TO_MAX_OUTPUT_TOKENS = {"command": 4096, "embed": 512}

    client: Client = field(kw_only=True)

    def count_tokens(self, text: str) -> int:
        return len(self.client.tokenize(text=text, model=self.model).tokens)

def count_tokens(self, text: str) -> int:
    return len(self.client.tokenize(text=text, model=self.model).tokens)

@define
class DummyTokenizer(BaseTokenizer):
    model: Optional[str] = field(default=None, kw_only=True)
    _max_input_tokens: int = field(init=False, default=0, kw_only=True, alias="max_input_tokens")
    _max_output_tokens: int = field(init=False, default=0, kw_only=True, alias="max_output_tokens")

    def count_tokens(self, text: str) -> int:
        raise DummyError(__class__.__name__, "count_tokens")

def count_tokens(self, text: str) -> int:
    raise DummyError(__class__.__name__, "count_tokens")

@define()
class GoogleTokenizer(BaseTokenizer):
    MODEL_PREFIXES_TO_MAX_INPUT_TOKENS = {"gemini-1.5-pro": 2097152, "gemini": 1048576}
    MODEL_PREFIXES_TO_MAX_OUTPUT_TOKENS = {"gemini": 8192}

    api_key: str = field(kw_only=True, metadata={"serializable": True})
    _client: Optional[GenerativeModel] = field(
        default=None, kw_only=True, alias="client", metadata={"serializable": False}
    )

    @lazy_property()
    def client(self) -> GenerativeModel:
        genai = import_optional_dependency("google.generativeai")
        genai.configure(api_key=self.api_key)

        return genai.GenerativeModel(self.model)

    def count_tokens(self, text: str) -> int:
        return self.client.count_tokens(text).total_tokens

@lazy_property()
def client(self) -> GenerativeModel:
    genai = import_optional_dependency("google.generativeai")
    genai.configure(api_key=self.api_key)

    return genai.GenerativeModel(self.model)

def count_tokens(self, text: str) -> int:
    return self.client.count_tokens(text).total_tokens

@define()
class GrokTokenizer(BaseTokenizer):
    # https://docs.x.ai/docs/models?cluster=us-east-1#model-constraints
    MODEL_PREFIXES_TO_MAX_INPUT_TOKENS = {
        "grok-2-vision": 32768,
        "grok-2": 131072,
        "grok-vision-beta": 8192,
        "grok-beta": 131072,
    }

    MODEL_PREFIXES_TO_MAX_OUTPUT_TOKENS = {"grok": 4096}

    base_url: str = field(default="https://api.x.ai", kw_only=True, metadata={"serializable": True})
    api_key: str = field(kw_only=True, default=None)
    headers: dict = field(
        default=Factory(lambda self: {"Authorization": f"Bearer {self.api_key}"}, takes_self=True), kw_only=True
    )

    def count_tokens(self, text: str) -> int:
        response = requests.post(
            urljoin(self.base_url, "/v1/tokenize-text"),
            headers=self.headers,
            json={"text": text, "model": self.model},
        )
        response.raise_for_status()
        return len(response.json()["token_ids"])

def count_tokens(self, text: str) -> int:
    response = requests.post(
        urljoin(self.base_url, "/v1/tokenize-text"),
        headers=self.headers,
        json={"text": text, "model": self.model},
    )
    response.raise_for_status()
    return len(response.json()["token_ids"])

@define()
class HuggingFaceTokenizer(BaseTokenizer):
    tokenizer: PreTrainedTokenizerBase = field(
        default=Factory(
            lambda self: import_optional_dependency("transformers").AutoTokenizer.from_pretrained(self.model),
            takes_self=True,
        ),
        kw_only=True,
    )
    _max_input_tokens: int = field(
        default=Factory(lambda self: self.tokenizer.model_max_length, takes_self=True),
        kw_only=True,
        alias="max_input_tokens",
    )
    _max_output_tokens: int = field(default=4096, kw_only=True, alias="max_output_tokens")

    def count_tokens(self, text: str) -> int:
        return len(self.tokenizer.encode(text))  # pyright: ignore[reportArgumentType]

def count_tokens(self, text: str) -> int:
    return len(self.tokenizer.encode(text))  # pyright: ignore[reportArgumentType]

@define()
class OpenAiTokenizer(BaseTokenizer):
    DEFAULT_OPENAI_GPT_3_COMPLETION_MODEL = "gpt-3.5-turbo-instruct"
    DEFAULT_OPENAI_GPT_3_CHAT_MODEL = "gpt-3.5-turbo"
    DEFAULT_OPENAI_GPT_4_MODEL = "gpt-4o"
    DEFAULT_ENCODING = "cl100k_base"
    DEFAULT_MAX_TOKENS = 2049
    DEFAULT_MAX_OUTPUT_TOKENS = 4096
    TOKEN_OFFSET = 8

    # https://platform.openai.com/docs/models/gpt-4-and-gpt-4-turbo
    MODEL_PREFIXES_TO_MAX_INPUT_TOKENS = {
        "gpt-4.1": 1_000_000,
        "gpt-4o": 128000,
        "gpt-4-1106": 128000,
        "gpt-4-32k": 32768,
        "gpt-4": 8192,
        "gpt-3.5-turbo-16k": 16384,
        "gpt-3.5-turbo": 4096,
        "gpt-35-turbo-16k": 16384,
        "gpt-35-turbo": 4096,
        "text-embedding-ada-002": 8191,
        "text-embedding-ada-001": 2046,
        "text-embedding-3-small": 8191,
        "text-embedding-3-large": 8191,
    }

    MODEL_PREFIXES_TO_MAX_OUTPUT_TOKENS = {"gpt": 4096}

    EMBEDDING_MODELS = [
        "text-embedding-ada-002",
        "text-embedding-ada-001",
        "text-embedding-3-small",
        "text-embedding-3-large",
    ]

    _max_input_tokens: int = field(
        kw_only=True,
        default=Factory(lambda self: self._default_max_input_tokens(), takes_self=True),
        alias="max_input_tokens",
    )
    _max_output_tokens: int = field(
        kw_only=True,
        default=Factory(lambda self: self._default_max_output_tokens(), takes_self=True),
        alias="max_output_tokens",
    )

    @property
    def encoding(self) -> tiktoken.Encoding:
        try:
            return tiktoken.encoding_for_model(self.model)
        except KeyError:
            return tiktoken.get_encoding(self.DEFAULT_ENCODING)

    def _default_max_input_tokens(self) -> int:
        tokens = next((v for k, v in self.MODEL_PREFIXES_TO_MAX_INPUT_TOKENS.items() if self.model.startswith(k)), None)
        offset = 0 if self.model in self.EMBEDDING_MODELS else self.TOKEN_OFFSET

        return (tokens or self.DEFAULT_MAX_TOKENS) - offset

    def _default_max_output_tokens(self) -> int:
        tokens = next(
            (v for k, v in self.MODEL_PREFIXES_TO_MAX_OUTPUT_TOKENS.items() if self.model.startswith(k)),
            None,
        )

        if tokens is None:
            return self.DEFAULT_MAX_OUTPUT_TOKENS
        return tokens

    def count_tokens(self, text: str | list[dict], model: Optional[str] = None) -> int:  # noqa: C901
        """Handles the special case of ChatML.

        Implementation adopted from the official OpenAI notebook:
        https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb.
        """
        if isinstance(text, list):
            model = model or self.model

            try:
                encoding = tiktoken.encoding_for_model(model)
            except KeyError:
                logging.warning("model not found. Using cl100k_base encoding.")

                encoding = tiktoken.get_encoding("cl100k_base")

            if model in {
                "gpt-3.5-turbo-0613",
                "gpt-3.5-turbo-16k-0613",
                "gpt-4-0314",
                "gpt-4-32k-0314",
                "gpt-4-0613",
                "gpt-4-32k-0613",
                "gpt-4o-2024-05-13",
            }:
                tokens_per_message = 3
                tokens_per_name = 1
            elif model == "gpt-3.5-turbo-0301":
                # every message follows <|start|>{role/name}\n{content}<|end|>\n
                tokens_per_message = 4
                # if there's a name, the role is omitted
                tokens_per_name = -1
            elif "gpt-3.5-turbo" in model or "gpt-35-turbo" in model:
                logging.info("gpt-3.5-turbo may update over time. Returning num tokens assuming gpt-3.5-turbo-0613.")
                return self.count_tokens(text, model="gpt-3.5-turbo-0613")
            elif "gpt-4o" in model:
                logging.info("gpt-4o may update over time. Returning num tokens assuming gpt-4o-2024-05-13.")
                return self.count_tokens(text, model="gpt-4o-2024-05-13")
            elif "gpt-4" in model:
                logging.info("gpt-4 may update over time. Returning num tokens assuming gpt-4-0613.")
                return self.count_tokens(text, model="gpt-4-0613")
            else:
                raise NotImplementedError(
                    f"token_count() is not implemented for model {model}. "
                    "See https://github.com/openai/openai-python/blob/main/chatml.md for "
                    "information on how messages are converted to tokens."
                )

            num_tokens = 0

            for message in text:
                num_tokens += tokens_per_message
                for key, value in message.items():
                    num_tokens += len(encoding.encode(value))
                    if key == "name":
                        num_tokens += tokens_per_name

            # every reply is primed with <|start|>assistant<|message|>
            num_tokens += 3

            return num_tokens
        return len(self.encoding.encode(text, allowed_special=set(self.stop_sequences)))

def _default_max_input_tokens(self) -> int:
    tokens = next((v for k, v in self.MODEL_PREFIXES_TO_MAX_INPUT_TOKENS.items() if self.model.startswith(k)), None)
    offset = 0 if self.model in self.EMBEDDING_MODELS else self.TOKEN_OFFSET

    return (tokens or self.DEFAULT_MAX_TOKENS) - offset

def _default_max_output_tokens(self) -> int:
    tokens = next(
        (v for k, v in self.MODEL_PREFIXES_TO_MAX_OUTPUT_TOKENS.items() if self.model.startswith(k)),
        None,
    )

    if tokens is None:
        return self.DEFAULT_MAX_OUTPUT_TOKENS
    return tokens

def count_tokens(self, text: str | list[dict], model: Optional[str] = None) -> int:  # noqa: C901
    """Handles the special case of ChatML.

    Implementation adopted from the official OpenAI notebook:
    https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb.
    """
    if isinstance(text, list):
        model = model or self.model

        try:
            encoding = tiktoken.encoding_for_model(model)
        except KeyError:
            logging.warning("model not found. Using cl100k_base encoding.")

            encoding = tiktoken.get_encoding("cl100k_base")

        if model in {
            "gpt-3.5-turbo-0613",
            "gpt-3.5-turbo-16k-0613",
            "gpt-4-0314",
            "gpt-4-32k-0314",
            "gpt-4-0613",
            "gpt-4-32k-0613",
            "gpt-4o-2024-05-13",
        }:
            tokens_per_message = 3
            tokens_per_name = 1
        elif model == "gpt-3.5-turbo-0301":
            # every message follows <|start|>{role/name}\n{content}<|end|>\n
            tokens_per_message = 4
            # if there's a name, the role is omitted
            tokens_per_name = -1
        elif "gpt-3.5-turbo" in model or "gpt-35-turbo" in model:
            logging.info("gpt-3.5-turbo may update over time. Returning num tokens assuming gpt-3.5-turbo-0613.")
            return self.count_tokens(text, model="gpt-3.5-turbo-0613")
        elif "gpt-4o" in model:
            logging.info("gpt-4o may update over time. Returning num tokens assuming gpt-4o-2024-05-13.")
            return self.count_tokens(text, model="gpt-4o-2024-05-13")
        elif "gpt-4" in model:
            logging.info("gpt-4 may update over time. Returning num tokens assuming gpt-4-0613.")
            return self.count_tokens(text, model="gpt-4-0613")
        else:
            raise NotImplementedError(
                f"token_count() is not implemented for model {model}. "
                "See https://github.com/openai/openai-python/blob/main/chatml.md for "
                "information on how messages are converted to tokens."
            )

        num_tokens = 0

        for message in text:
            num_tokens += tokens_per_message
            for key, value in message.items():
                num_tokens += len(encoding.encode(value))
                if key == "name":
                    num_tokens += tokens_per_name

        # every reply is primed with <|start|>assistant<|message|>
        num_tokens += 3

        return num_tokens
    return len(self.encoding.encode(text, allowed_special=set(self.stop_sequences)))

@define()
class SimpleTokenizer(BaseTokenizer):
    model: str | None = field(init=False, default=None, kw_only=True)
    characters_per_token: int = field(kw_only=True)

    def count_tokens(self, text: str) -> int:
        return (len(text) + self.characters_per_token - 1) // self.characters_per_token

def count_tokens(self, text: str) -> int:
    return (len(text) + self.characters_per_token - 1) // self.characters_per_token

@define()
class VoyageAiTokenizer(BaseTokenizer):
    MODEL_PREFIXES_TO_MAX_INPUT_TOKENS = {
        "voyage-3": 32000,
        "voyage-large-2": 16000,
        "voyage-code-2": 16000,
        "voyage-2": 4000,
        "voyage-lite-02-instruct": 4000,
    }
    MODEL_PREFIXES_TO_MAX_OUTPUT_TOKENS = {"voyage": 0}

    api_key: Optional[str] = field(default=None, kw_only=True, metadata={"serializable": False})
    client: Client = field(
        default=Factory(
            lambda self: import_optional_dependency("voyageai").Client(api_key=self.api_key),
            takes_self=True,
        ),
        kw_only=True,
    )

    def count_tokens(self, text: str) -> int:
        return self.client.count_tokens([text], model=self.model)

def count_tokens(self, text: str) -> int:
    return self.client.count_tokens([text], model=self.model)