Token 计数器

演示如何使用 tiktoken 估算文本的 token 数量并估算调用成本。 当 tiktoken 不可用时，自动回退到基于启发式的近似估算，保证示例可离线运行。 对应教程：第1卷《AI 基础》第3章 Token 与上下文窗口。

文件路径: examples/volume-01-ai-fundamentals\token_counter.py

---

源代码

"""Token 计算与成本估算示例。

演示如何使用 ``tiktoken`` 估算文本的 token 数量并估算调用成本。
当 ``tiktoken`` 不可用时，自动回退到基于启发式的近似估算，保证示例可离线运行。

对应教程：第1卷《AI 基础》第3章 Token 与上下文窗口。
"""

from __future__ import annotations

from dataclasses import dataclass

try:  # tiktoken 在 requirements.txt 中，但这里做防御性回退
    import tiktoken

    _HAS_TIKTOKEN = True
except Exception:  # pragma: no cover - tiktoken 通常可用  # noqa: BLE001
    tiktoken = None  # type: ignore[assignment]
    _HAS_TIKTOKEN = False


# 常见模型上下文窗口与计价（美元 / 1K token），仅作教学示例
MODEL_PROFILES: dict[str, dict[str, float]] = {
    "gpt-4o": {"context_window": 128_000, "input_price": 0.0025, "output_price": 0.01},
    "gpt-4o-mini": {"context_window": 128_000, "input_price": 0.00015, "output_price": 0.0006},
    "gpt-3.5-turbo": {"context_window": 16_385, "input_price": 0.0005, "output_price": 0.0015},
}


@dataclass(frozen=True)
class TokenUsage:
    """Token 用量统计。

    Attributes:
        prompt_tokens: 输入 token 数。
        completion_tokens: 输出 token 数。
        total_tokens: 总 token 数。
    """

    prompt_tokens: int
    completion_tokens: int

    @property
    def total_tokens(self) -> int:
        return self.prompt_tokens + self.completion_tokens


@dataclass(frozen=True)
class CostEstimate:
    """成本估算结果。

    Attributes:
        usage: token 用量。
        model: 模型名称。
        input_cost: 输入成本（美元）。
        output_cost: 输出成本（美元）。
    """

    usage: TokenUsage
    model: str
    input_cost: float
    output_cost: float

    @property
    def total_cost(self) -> float:
        return self.input_cost + self.output_cost


def estimate_tokens_heuristic(text: str) -> int:
    """用启发式规则近似估算 token 数。

    规则：
        - 英文/ASCII：约 1 token ≈ 4 字符。
        - 中文字符：约每个字 1.5 token。
        - 至少为 1，避免空文本返回 0 造成误解。

    Args:
        text: 待估算的文本。

    Returns:
        估算的 token 数。
    """
    if not text:
        return 0
    cjk = sum(1 for ch in text if "\u4e00" <= ch <= "\u9fff")
    other = len(text) - cjk
    return max(1, int(cjk * 1.5 + other / 4))


def count_tokens(text: str, model: str = "gpt-4o") -> int:
    """计算文本在指定模型下的 token 数。

    优先使用 ``tiktoken``；不可用时回退到启发式估算。

    Args:
        text: 待计算的文本。
        model: 模型名称，用于选择分词器。

    Returns:
        token 数量。
    """
    if _HAS_TIKTOKEN:
        try:
            encoding = tiktoken.encoding_for_model(model)
            return len(encoding.encode(text))
        except Exception:  # noqa: BLE001 - 未知模型名时尝试通用编码
            try:
                encoding = tiktoken.get_encoding("cl100k_base")
                return len(encoding.encode(text))
            except Exception:  # noqa: BLE001 - pragma: no cover
                pass
    return estimate_tokens_heuristic(text)


def count_messages_tokens(messages: list[dict[str, str]], model: str = "gpt-4o") -> int:
    """估算一组聊天消息的 token 数。

    每条消息附带约 4 个固定 token 的开销（角色标记等），最后再加 2 个 token 收尾。

    Args:
        messages: 消息列表，每项形如 ``{"role": ..., "content": ...}``。
        model: 模型名称。

    Returns:
        总 token 数。
    """
    total = 0
    for msg in messages:
        total += 4  # 每条消息固定开销
        for key in ("role", "content", "name"):
            if key in msg and isinstance(msg[key], str):
                total += count_tokens(msg[key], model)
    total += 2  # 收尾 token
    return total


def estimate_cost(
    prompt: str | list[dict[str, str]],
    completion: str = "",
    model: str = "gpt-4o",
) -> CostEstimate:
    """估算一次调用的 token 用量与成本。

    Args:
        prompt: 输入文本或消息列表。
        completion: 输出文本。
        model: 模型名称，需在 ``MODEL_PROFILES`` 中。

    Returns:
        ``CostEstimate`` 估算结果。

    Raises:
        KeyError: 模型不在 ``MODEL_PROFILES`` 中时。
    """
    if model not in MODEL_PROFILES:
        raise KeyError(f"未知模型：{model}，可用：{list(MODEL_PROFILES)}")
    profile = MODEL_PROFILES[model]

    if isinstance(prompt, str):
        prompt_tokens = count_tokens(prompt, model)
    else:
        prompt_tokens = count_messages_tokens(prompt, model)
    completion_tokens = count_tokens(completion, model)

    usage = TokenUsage(prompt_tokens=prompt_tokens, completion_tokens=completion_tokens)
    input_cost = usage.prompt_tokens / 1000 * profile["input_price"]
    output_cost = usage.completion_tokens / 1000 * profile["output_price"]
    return CostEstimate(usage=usage, model=model, input_cost=input_cost, output_cost=output_cost)


def fits_in_context(prompt: str, model: str = "gpt-4o", reserved_output: int = 1024) -> bool:
    """判断文本是否在模型上下文窗口内。

    Args:
        prompt: 输入文本。
        model: 模型名称。
        reserved_output: 为输出预留的 token 数。

    Returns:
        若 ``输入 token + 预留输出 <= 上下文窗口`` 返回 True。
    """
    profile = MODEL_PROFILES.get(model, {"context_window": 4096})
    used = count_tokens(prompt, model) + reserved_output
    return used <= profile["context_window"]


def main() -> None:
    """演示 token 计算与成本估算。"""
    text = "AI Agent 是能够感知环境、自主决策并执行任务的智能体。"
    print(f"文本: {text}")
    print(f"字符数: {len(text)}")
    print(f"token 数 (tiktoken/heuristic): {count_tokens(text)}")
    print(f"启发式估算: {estimate_tokens_heuristic(text)}\n")

    for model_name in MODEL_PROFILES:
        est = estimate_cost(text, "好的，我已经理解了。", model=model_name)
        print(
            f"[{model_name}] 输入={est.usage.prompt_tokens} 输出={est.usage.completion_tokens} "
            f"合计=${est.total_cost:.6f}"
        )


if __name__ == "__main__":
    main()