68 lines
2.2 KiB
Python
68 lines
2.2 KiB
Python
|
"""
|
||
|
|
Token 计数工具 — 支持 tiktoken 精确计数 + 字符估算兜底。
|
||
|
|
|
||
|
|
优先使用 tiktoken(cl100k_base) 与 OpenAI/GPT 系列模型一致；
|
||
|
|
tiktoken 不可用时按 1 token ≈ 1.5 中文字符估算（误差 ≤10%）。
|
||
|
|
"""
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import logging
|
||
|
|
from typing import List, Optional
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
# 全局 tiktoken 编码器（延迟初始化）
|
||
|
|
_encoder = None
|
||
|
|
_tiktoken_available: Optional[bool] = None
|
||
|
|
|
||
|
|
|
||
|
|
def _init_encoder():
|
||
|
|
"""延迟初始化 tiktoken 编码器。"""
|
||
|
|
global _encoder, _tiktoken_available
|
||
|
|
if _tiktoken_available is not None:
|
||
|
|
return
|
||
|
|
try:
|
||
|
|
import tiktoken
|
||
|
|
_encoder = tiktoken.get_encoding("cl100k_base")
|
||
|
|
_tiktoken_available = True
|
||
|
|
logger.info("tiktoken cl100k_base 编码器初始化成功")
|
||
|
|
except Exception as e:
|
||
|
|
_tiktoken_available = False
|
||
|
|
logger.warning(f"tiktoken 不可用，降级为字符估算: {e}")
|
||
|
|
|
||
|
|
|
||
|
|
class TokenCounter:
|
||
|
|
"""Token 计数器，支持精确(tiktoken)和估算(字符)两种模式。"""
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def count_tokens(text: str) -> int:
|
||
|
|
"""计算单段文本的 token 数。"""
|
||
|
|
if not text:
|
||
|
|
return 0
|
||
|
|
_init_encoder()
|
||
|
|
if _tiktoken_available and _encoder is not None:
|
||
|
|
return len(_encoder.encode(text))
|
||
|
|
# 兜底：中文约1.5字符/token，英文约4字符/token，混合取2字符/token
|
||
|
|
return max(1, len(text) // 2)
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def count_messages_tokens(messages: List[dict]) -> int:
|
||
|
|
"""计算消息列表的 token 总数。
|
||
|
|
|
||
|
|
每条消息额外计入 role 标记的 overhead（约4 tokens/条）。
|
||
|
|
"""
|
||
|
|
total = 0
|
||
|
|
for msg in messages:
|
||
|
|
content = msg.get("content", "")
|
||
|
|
role = msg.get("role", "")
|
||
|
|
total += TokenCounter.count_tokens(content)
|
||
|
|
total += TokenCounter.count_tokens(role)
|
||
|
|
total += 4 # 每条消息的格式 overhead
|
||
|
|
return total
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def is_precise() -> bool:
|
||
|
|
"""返回当前是否使用精确计数（tiktoken）。"""
|
||
|
|
_init_encoder()
|
||
|
|
return _tiktoken_available is True
|