mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-01 17:00:21 +08:00
refactor: 移除 Python 后端源码,全面迁移至 Rust gateway 架构
- 删除全部 Python 源码 (src/) 及 Alembic 迁移脚本,归档至 _deprecated_py_src/ - 重构 Rust gateway ai_pipeline: 拆分 planner/finalize 模块,新增 contracts/adaptation 层 - 重组 handlers 模块为 admin/public/proxy/internal/shared 子模块结构 - 新增 executor 模块,引入 Rust 原生数据库迁移 (aether-data/migrations) - 简化 CI/Docker 构建流程,移除 base image 二级构建,统一为单一 app image - 移除 Python 相关基础设施文件 (entrypoint.sh, gunicorn_conf.py, Dockerfile.base)
This commit is contained in:
9
_deprecated_py_src/plugins/token/__init__.py
Normal file
9
_deprecated_py_src/plugins/token/__init__.py
Normal file
@@ -0,0 +1,9 @@
|
||||
"""
|
||||
Token计数插件
|
||||
"""
|
||||
|
||||
from .base import TokenCounterPlugin, TokenUsage
|
||||
from .claude_counter import ClaudeTokenCounterPlugin
|
||||
from .tiktoken_counter import TiktokenCounterPlugin
|
||||
|
||||
__all__ = ["TokenCounterPlugin", "TokenUsage", "TiktokenCounterPlugin", "ClaudeTokenCounterPlugin"]
|
||||
170
_deprecated_py_src/plugins/token/base.py
Normal file
170
_deprecated_py_src/plugins/token/base.py
Normal file
@@ -0,0 +1,170 @@
|
||||
"""
|
||||
Token计数插件基类
|
||||
定义Token计数的接口
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from src.plugins.common import BasePlugin
|
||||
|
||||
|
||||
@dataclass
|
||||
class TokenUsage:
|
||||
"""令牌使用情况"""
|
||||
|
||||
input_tokens: int = 0
|
||||
output_tokens: int = 0
|
||||
total_tokens: int = 0
|
||||
cache_read_tokens: int = 0 # Claude缓存读取
|
||||
cache_write_tokens: int = 0 # Claude缓存写入
|
||||
reasoning_tokens: int = 0 # OpenAI o1推理令牌
|
||||
|
||||
def __add__(self, other: TokenUsage) -> TokenUsage:
|
||||
"""令牌使用相加"""
|
||||
return TokenUsage(
|
||||
input_tokens=self.input_tokens + other.input_tokens,
|
||||
output_tokens=self.output_tokens + other.output_tokens,
|
||||
total_tokens=self.total_tokens + other.total_tokens,
|
||||
cache_read_tokens=self.cache_read_tokens + other.cache_read_tokens,
|
||||
cache_write_tokens=self.cache_write_tokens + other.cache_write_tokens,
|
||||
reasoning_tokens=self.reasoning_tokens + other.reasoning_tokens,
|
||||
)
|
||||
|
||||
def to_dict(self) -> dict[str, int]:
|
||||
"""转换为字典"""
|
||||
return {
|
||||
"input_tokens": self.input_tokens,
|
||||
"output_tokens": self.output_tokens,
|
||||
"total_tokens": self.total_tokens,
|
||||
"cache_read_tokens": self.cache_read_tokens,
|
||||
"cache_write_tokens": self.cache_write_tokens,
|
||||
"reasoning_tokens": self.reasoning_tokens,
|
||||
}
|
||||
|
||||
|
||||
class TokenCounterPlugin(BasePlugin):
|
||||
"""
|
||||
Token计数插件基类
|
||||
支持不同模型的Token计数
|
||||
"""
|
||||
|
||||
def __init__(self, name: str = "token_counter", config: dict[str, Any] | None = None):
|
||||
# 调用父类初始化,设置metadata
|
||||
super().__init__(
|
||||
name=name, config=config, description="Token Counter Plugin", version="1.0.0"
|
||||
)
|
||||
|
||||
self.supported_models = self.config.get("supported_models", [])
|
||||
self.default_model = self.config.get("default_model")
|
||||
|
||||
@abstractmethod
|
||||
def supports_model(self, model: str) -> bool:
|
||||
"""检查是否支持指定模型"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def count_tokens(self, text: str, model: str | None = None) -> int:
|
||||
"""计算文本的Token数量"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def count_messages(self, messages: list[dict[str, Any]], model: str | None = None) -> int:
|
||||
"""计算消息列表的Token数量"""
|
||||
pass
|
||||
|
||||
async def count_request(self, request: dict[str, Any], model: str | None = None) -> int:
|
||||
"""计算请求的Token数量"""
|
||||
model = model or request.get("model") or self.default_model
|
||||
messages = request.get("messages", [])
|
||||
return await self.count_messages(messages, model)
|
||||
|
||||
async def count_response(
|
||||
self, response: dict[str, Any], model: str | None = None
|
||||
) -> TokenUsage:
|
||||
"""从响应中提取Token使用情况"""
|
||||
usage = response.get("usage", {})
|
||||
|
||||
# OpenAI格式
|
||||
if "prompt_tokens" in usage:
|
||||
return TokenUsage(
|
||||
input_tokens=usage.get("prompt_tokens", 0),
|
||||
output_tokens=usage.get("completion_tokens", 0),
|
||||
total_tokens=usage.get("total_tokens", 0),
|
||||
reasoning_tokens=usage.get("completion_tokens_details", {}).get(
|
||||
"reasoning_tokens", 0
|
||||
),
|
||||
)
|
||||
|
||||
# Claude格式
|
||||
elif "input_tokens" in usage:
|
||||
return TokenUsage(
|
||||
input_tokens=usage.get("input_tokens", 0),
|
||||
output_tokens=usage.get("output_tokens", 0),
|
||||
total_tokens=usage.get("input_tokens", 0) + usage.get("output_tokens", 0),
|
||||
cache_read_tokens=usage.get("cache_read_input_tokens", 0),
|
||||
cache_write_tokens=usage.get("cache_creation_input_tokens", 0),
|
||||
)
|
||||
|
||||
return TokenUsage()
|
||||
|
||||
async def estimate_cost(
|
||||
self, usage: TokenUsage, model: str, provider: str | None = None
|
||||
) -> dict[str, float]:
|
||||
"""估算使用成本"""
|
||||
# 默认价格表(每1M tokens的价格)
|
||||
pricing = self.config.get("pricing", {})
|
||||
|
||||
# 获取模型价格
|
||||
model_pricing = pricing.get(model, {})
|
||||
if not model_pricing:
|
||||
# 尝试使用前缀匹配
|
||||
for model_prefix, price_info in pricing.items():
|
||||
if model.startswith(model_prefix):
|
||||
model_pricing = price_info
|
||||
break
|
||||
|
||||
if not model_pricing:
|
||||
return {"error": "No pricing information available"}
|
||||
|
||||
# 计算成本
|
||||
input_cost = (usage.input_tokens / 1_000_000) * model_pricing.get("input", 0)
|
||||
output_cost = (usage.output_tokens / 1_000_000) * model_pricing.get("output", 0)
|
||||
|
||||
# 缓存成本(Claude特有)
|
||||
cache_read_cost = (usage.cache_read_tokens / 1_000_000) * model_pricing.get("cache_read", 0)
|
||||
cache_write_cost = (usage.cache_write_tokens / 1_000_000) * model_pricing.get(
|
||||
"cache_write", 0
|
||||
)
|
||||
|
||||
# 推理成本(OpenAI o1特有)
|
||||
reasoning_cost = (usage.reasoning_tokens / 1_000_000) * model_pricing.get("reasoning", 0)
|
||||
|
||||
total_cost = input_cost + output_cost + cache_read_cost + cache_write_cost + reasoning_cost
|
||||
|
||||
return {
|
||||
"input_cost": round(input_cost, 6),
|
||||
"output_cost": round(output_cost, 6),
|
||||
"cache_read_cost": round(cache_read_cost, 6),
|
||||
"cache_write_cost": round(cache_write_cost, 6),
|
||||
"reasoning_cost": round(reasoning_cost, 6),
|
||||
"total_cost": round(total_cost, 6),
|
||||
"currency": "USD",
|
||||
}
|
||||
|
||||
@abstractmethod
|
||||
async def get_model_info(self, model: str) -> dict[str, Any]:
|
||||
"""获取模型信息"""
|
||||
pass
|
||||
|
||||
async def get_stats(self) -> dict[str, Any]:
|
||||
"""获取统计信息"""
|
||||
return {
|
||||
"type": self.name,
|
||||
"enabled": self.enabled,
|
||||
"supported_models": self.supported_models,
|
||||
"default_model": self.default_model,
|
||||
}
|
||||
274
_deprecated_py_src/plugins/token/claude_counter.py
Normal file
274
_deprecated_py_src/plugins/token/claude_counter.py
Normal file
@@ -0,0 +1,274 @@
|
||||
"""
|
||||
Claude Token计数插件
|
||||
专门为Claude模型设计的Token计数器
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
from .base import TokenCounterPlugin
|
||||
|
||||
_CJK_PATTERN = re.compile(r"[\u4e00-\u9fff\u3040-\u30ff\uac00-\ud7af]")
|
||||
|
||||
|
||||
class ClaudeTokenCounterPlugin(TokenCounterPlugin):
|
||||
"""
|
||||
Claude专用Token计数插件
|
||||
使用简化的估算方法
|
||||
"""
|
||||
|
||||
# Claude模型信息
|
||||
CLAUDE_MODELS = {
|
||||
"claude-3-5-sonnet-20241022": {
|
||||
"max_tokens": 200000,
|
||||
"max_output": 8192,
|
||||
"chars_per_token": 3.5, # 平均字符/token比例
|
||||
},
|
||||
"claude-3-5-haiku-20241022": {
|
||||
"max_tokens": 200000,
|
||||
"max_output": 8192,
|
||||
"chars_per_token": 3.5,
|
||||
},
|
||||
"claude-3-opus-20240229": {
|
||||
"max_tokens": 200000,
|
||||
"max_output": 4096,
|
||||
"chars_per_token": 3.5,
|
||||
},
|
||||
"claude-3-sonnet-20240229": {
|
||||
"max_tokens": 200000,
|
||||
"max_output": 4096,
|
||||
"chars_per_token": 3.5,
|
||||
},
|
||||
"claude-3-haiku-20240307": {
|
||||
"max_tokens": 200000,
|
||||
"max_output": 4096,
|
||||
"chars_per_token": 3.5,
|
||||
},
|
||||
# 旧版模型
|
||||
"claude-2.1": {
|
||||
"max_tokens": 100000,
|
||||
"max_output": 4096,
|
||||
"chars_per_token": 4,
|
||||
},
|
||||
"claude-2.0": {
|
||||
"max_tokens": 100000,
|
||||
"max_output": 4096,
|
||||
"chars_per_token": 4,
|
||||
},
|
||||
"claude-instant-1.2": {
|
||||
"max_tokens": 100000,
|
||||
"max_output": 4096,
|
||||
"chars_per_token": 4,
|
||||
},
|
||||
}
|
||||
|
||||
def __init__(self, name: str = "claude", config: dict[str, Any] | None = None):
|
||||
super().__init__(name, config)
|
||||
|
||||
# 价格表(每1M tokens的价格 USD)
|
||||
default_pricing = {
|
||||
"claude-3-5-sonnet": {
|
||||
"input": 3,
|
||||
"output": 15,
|
||||
"cache_write": 3.75, # 缓存写入
|
||||
"cache_read": 0.30, # 缓存读取
|
||||
},
|
||||
"claude-3-5-haiku": {
|
||||
"input": 0.8,
|
||||
"output": 4,
|
||||
"cache_write": 1,
|
||||
"cache_read": 0.08,
|
||||
},
|
||||
"claude-3-opus": {
|
||||
"input": 15,
|
||||
"output": 75,
|
||||
"cache_write": 18.75,
|
||||
"cache_read": 1.50,
|
||||
},
|
||||
"claude-3-sonnet": {
|
||||
"input": 3,
|
||||
"output": 15,
|
||||
"cache_write": 3.75,
|
||||
"cache_read": 0.30,
|
||||
},
|
||||
"claude-3-haiku": {
|
||||
"input": 0.25,
|
||||
"output": 1.25,
|
||||
"cache_write": 0.30,
|
||||
"cache_read": 0.03,
|
||||
},
|
||||
"claude-2.1": {
|
||||
"input": 8,
|
||||
"output": 24,
|
||||
},
|
||||
"claude-2.0": {
|
||||
"input": 8,
|
||||
"output": 24,
|
||||
},
|
||||
"claude-instant": {
|
||||
"input": 0.8,
|
||||
"output": 2.4,
|
||||
},
|
||||
}
|
||||
self.config["pricing"] = (
|
||||
config.get("pricing", default_pricing) if config else default_pricing
|
||||
)
|
||||
|
||||
def supports_model(self, model: str) -> bool:
|
||||
"""检查是否支持指定模型"""
|
||||
# 支持所有Claude模型
|
||||
return "claude" in model.lower()
|
||||
|
||||
def _estimate_tokens_from_text(self, text: str, model: str) -> int:
|
||||
"""从文本估算Token数量"""
|
||||
# 获取模型信息
|
||||
model_info = None
|
||||
for model_name, info in self.CLAUDE_MODELS.items():
|
||||
if model.startswith(model_name.split("-20")[0]): # 匹配基本名称
|
||||
model_info = info
|
||||
break
|
||||
|
||||
if not model_info:
|
||||
# 默认值
|
||||
model_info = {"chars_per_token": 3.5}
|
||||
|
||||
# 基本估算
|
||||
chars_per_token = model_info["chars_per_token"]
|
||||
|
||||
# 考虑不同语言的特点
|
||||
# 检测是否包含中文/日文/韩文
|
||||
cjk_count = len(_CJK_PATTERN.findall(text))
|
||||
|
||||
if cjk_count > len(text) * 0.3: # 超过30%是CJK字符
|
||||
# CJK字符通常每个字符1-2个token
|
||||
return int(len(text) / 1.5)
|
||||
else:
|
||||
# 英文和其他语言
|
||||
# 考虑空格和标点
|
||||
word_count = len(text.split())
|
||||
# 平均每个单词1.3个token
|
||||
token_by_words = int(word_count * 1.3)
|
||||
# 平均每个字符chars_per_token
|
||||
token_by_chars = int(len(text) / chars_per_token)
|
||||
# 取两者的平均
|
||||
return (token_by_words + token_by_chars) // 2
|
||||
|
||||
async def count_tokens(self, text: str, model: str | None = None) -> int:
|
||||
"""计算文本的Token数量"""
|
||||
if not self.enabled:
|
||||
return 0
|
||||
|
||||
model = model or self.default_model or "claude-3-5-sonnet-20241022"
|
||||
return self._estimate_tokens_from_text(text, model)
|
||||
|
||||
async def count_messages(self, messages: list[dict[str, Any]], model: str | None = None) -> int:
|
||||
"""计算消息列表的Token数量"""
|
||||
if not self.enabled:
|
||||
return 0
|
||||
|
||||
model = model or self.default_model or "claude-3-5-sonnet-20241022"
|
||||
total_tokens = 0
|
||||
|
||||
for message in messages:
|
||||
# 角色token(约3 tokens)
|
||||
total_tokens += 3
|
||||
|
||||
# 内容token
|
||||
content = message.get("content")
|
||||
if content:
|
||||
if isinstance(content, str):
|
||||
total_tokens += self._estimate_tokens_from_text(content, model)
|
||||
elif isinstance(content, list):
|
||||
# 处理多模态内容
|
||||
for item in content:
|
||||
if item.get("type") == "text":
|
||||
text = item.get("text", "")
|
||||
total_tokens += self._estimate_tokens_from_text(text, model)
|
||||
elif item.get("type") == "image":
|
||||
# Claude图像处理
|
||||
# 基础: 1,600 tokens
|
||||
# 每个256x256的tile: 280 tokens
|
||||
# 简化估算
|
||||
total_tokens += 2000 # 平均估算
|
||||
elif item.get("type") == "tool_use":
|
||||
# 工具使用
|
||||
tool_name = item.get("name", "")
|
||||
tool_input = item.get("input", {})
|
||||
total_tokens += self._estimate_tokens_from_text(tool_name, model)
|
||||
total_tokens += self._estimate_tokens_from_text(
|
||||
json.dumps(tool_input), model
|
||||
)
|
||||
elif item.get("type") == "tool_result":
|
||||
# 工具结果
|
||||
tool_content = item.get("content", "")
|
||||
if isinstance(tool_content, str):
|
||||
total_tokens += self._estimate_tokens_from_text(tool_content, model)
|
||||
|
||||
# 添加系统提示的token(如果有)
|
||||
if messages and messages[0].get("role") == "system":
|
||||
# 系统提示通常会有额外的开销
|
||||
total_tokens += 10
|
||||
|
||||
return total_tokens
|
||||
|
||||
async def count_request(self, request: dict[str, Any], model: str | None = None) -> int:
|
||||
"""计算请求的Token数量"""
|
||||
model = model or request.get("model") or self.default_model
|
||||
messages = request.get("messages", [])
|
||||
total = await self.count_messages(messages, model)
|
||||
|
||||
# 考虑系统提示
|
||||
system = request.get("system")
|
||||
if system:
|
||||
total += self._estimate_tokens_from_text(system, model)
|
||||
total += 5 # 系统提示的额外开销
|
||||
|
||||
return total
|
||||
|
||||
async def get_model_info(self, model: str) -> dict[str, Any]:
|
||||
"""获取模型信息"""
|
||||
info = {"model": model, "supported": self.supports_model(model)}
|
||||
|
||||
if self.supports_model(model):
|
||||
# 查找匹配的模型信息
|
||||
model_info = None
|
||||
for model_name, m_info in self.CLAUDE_MODELS.items():
|
||||
if model.startswith(model_name.split("-20")[0]):
|
||||
model_info = m_info
|
||||
info["model_name"] = model_name
|
||||
break
|
||||
|
||||
if model_info:
|
||||
info.update(
|
||||
{
|
||||
"max_tokens": model_info["max_tokens"],
|
||||
"max_output": model_info["max_output"],
|
||||
"chars_per_token": model_info["chars_per_token"],
|
||||
"supports_vision": "claude-3" in model,
|
||||
"supports_tools": "claude-3" in model,
|
||||
"supports_cache": "claude-3" in model,
|
||||
}
|
||||
)
|
||||
|
||||
# 添加价格信息
|
||||
pricing = self.config.get("pricing", {})
|
||||
for price_key in pricing:
|
||||
if model.startswith(price_key):
|
||||
info["pricing"] = pricing[price_key]
|
||||
break
|
||||
|
||||
return info
|
||||
|
||||
async def get_stats(self) -> dict[str, Any]:
|
||||
"""获取统计信息"""
|
||||
stats = await super().get_stats()
|
||||
stats.update(
|
||||
{
|
||||
"estimation_method": "character_based",
|
||||
"supported_models_count": len(self.CLAUDE_MODELS),
|
||||
}
|
||||
)
|
||||
return stats
|
||||
286
_deprecated_py_src/plugins/token/tiktoken_counter.py
Normal file
286
_deprecated_py_src/plugins/token/tiktoken_counter.py
Normal file
@@ -0,0 +1,286 @@
|
||||
"""
|
||||
Tiktoken Token计数插件
|
||||
支持OpenAI和其他使用tiktoken的模型
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from functools import lru_cache
|
||||
from typing import Any
|
||||
|
||||
from src.core.logger import logger
|
||||
|
||||
from .base import TokenCounterPlugin
|
||||
|
||||
# 尝试导入tiktoken
|
||||
try:
|
||||
import tiktoken
|
||||
|
||||
TIKTOKEN_AVAILABLE = True
|
||||
except ImportError: # pragma: no cover
|
||||
TIKTOKEN_AVAILABLE = False
|
||||
tiktoken = None
|
||||
|
||||
|
||||
@lru_cache(maxsize=4)
|
||||
def _get_encoder_cached(model: str) -> Any:
|
||||
"""全局编码器缓存。
|
||||
|
||||
目的:避免在多实例/多请求场景下重复初始化 tiktoken 编码器。
|
||||
实际只有 cl100k_base / o200k_base / p50k_base 等少数几种编码,4 个足够。
|
||||
"""
|
||||
if not TIKTOKEN_AVAILABLE:
|
||||
raise RuntimeError("tiktoken not installed")
|
||||
|
||||
mapping = TiktokenCounterPlugin.MODEL_ENCODINGS
|
||||
|
||||
# 1) 完全匹配
|
||||
if model in mapping:
|
||||
return tiktoken.get_encoding(mapping[model])
|
||||
|
||||
# 2) 前缀匹配(按前缀长度从长到短,避免短前缀抢先匹配)
|
||||
for model_prefix, enc_name in TiktokenCounterPlugin.MODEL_ENCODINGS_PREFIXES:
|
||||
if model.startswith(model_prefix):
|
||||
return tiktoken.get_encoding(enc_name)
|
||||
|
||||
# 3) 尝试使用模型名称
|
||||
try:
|
||||
return tiktoken.encoding_for_model(model)
|
||||
except Exception:
|
||||
# 默认使用 cl100k_base
|
||||
return tiktoken.get_encoding("cl100k_base")
|
||||
|
||||
|
||||
class TiktokenCounterPlugin(TokenCounterPlugin):
|
||||
"""
|
||||
使用tiktoken库计算Token数量
|
||||
支持OpenAI模型和其他兼容模型
|
||||
"""
|
||||
|
||||
# 模型编码映射
|
||||
MODEL_ENCODINGS = {
|
||||
# GPT-4 系列
|
||||
"gpt-4": "cl100k_base",
|
||||
"gpt-4-32k": "cl100k_base",
|
||||
"gpt-4-turbo": "cl100k_base",
|
||||
"gpt-4-turbo-preview": "cl100k_base",
|
||||
"gpt-4o": "o200k_base",
|
||||
"gpt-4o-mini": "o200k_base",
|
||||
# GPT-3.5 系列
|
||||
"gpt-3.5-turbo": "cl100k_base",
|
||||
"gpt-3.5-turbo-16k": "cl100k_base",
|
||||
# 旧模型
|
||||
"text-davinci-003": "p50k_base",
|
||||
"text-davinci-002": "p50k_base",
|
||||
"code-davinci-002": "p50k_base",
|
||||
# Embeddings
|
||||
"text-embedding-ada-002": "cl100k_base",
|
||||
"text-embedding-3-small": "cl100k_base",
|
||||
"text-embedding-3-large": "cl100k_base",
|
||||
}
|
||||
|
||||
# 前缀匹配顺序(从长到短)
|
||||
MODEL_ENCODINGS_PREFIXES = sorted(
|
||||
MODEL_ENCODINGS.items(),
|
||||
key=lambda kv: len(kv[0]),
|
||||
reverse=True,
|
||||
)
|
||||
|
||||
# 每个消息的额外Token数
|
||||
MESSAGE_OVERHEAD = {
|
||||
"gpt-3.5-turbo": 4, # 每条消息
|
||||
"gpt-4": 3,
|
||||
"gpt-4-turbo": 3,
|
||||
"gpt-4o": 3,
|
||||
"gpt-4o-mini": 3,
|
||||
}
|
||||
|
||||
def __init__(self, name: str = "tiktoken", config: dict[str, Any] | None = None):
|
||||
super().__init__(name, config)
|
||||
|
||||
if not TIKTOKEN_AVAILABLE:
|
||||
self.enabled = False
|
||||
logger.warning("tiktoken not installed, plugin disabled")
|
||||
return
|
||||
|
||||
# 缓存编码器
|
||||
self._encoders = {}
|
||||
|
||||
# 价格表(每1M tokens的价格 USD)
|
||||
default_pricing = {
|
||||
"gpt-4o": {"input": 2.5, "output": 10},
|
||||
"gpt-4o-mini": {"input": 0.15, "output": 0.6},
|
||||
"gpt-4-turbo": {"input": 10, "output": 30},
|
||||
"gpt-4": {"input": 30, "output": 60},
|
||||
"gpt-3.5-turbo": {"input": 0.5, "output": 1.5},
|
||||
"o1-preview": {"input": 15, "output": 60, "reasoning": 60},
|
||||
"o1-mini": {"input": 3, "output": 12, "reasoning": 12},
|
||||
}
|
||||
self.config["pricing"] = (
|
||||
config.get("pricing", default_pricing) if config else default_pricing
|
||||
)
|
||||
|
||||
def _get_encoder(self, model: str) -> Any:
|
||||
"""获取模型的编码器(全局缓存)"""
|
||||
return _get_encoder_cached(model)
|
||||
|
||||
def supports_model(self, model: str) -> bool:
|
||||
"""检查是否支持指定模型"""
|
||||
# 支持所有OpenAI模型和一些兼容模型
|
||||
openai_models = [
|
||||
"gpt-4",
|
||||
"gpt-3.5",
|
||||
"text-davinci",
|
||||
"text-embedding",
|
||||
"code-davinci",
|
||||
"o1",
|
||||
]
|
||||
return any(model.startswith(prefix) for prefix in openai_models)
|
||||
|
||||
async def count_tokens(self, text: str, model: str | None = None) -> int:
|
||||
"""计算文本的Token数量"""
|
||||
if not self.enabled:
|
||||
return 0
|
||||
|
||||
model = model or self.default_model or "gpt-3.5-turbo"
|
||||
encoder = self._get_encoder(model)
|
||||
|
||||
try:
|
||||
tokens = encoder.encode(text)
|
||||
return len(tokens)
|
||||
except Exception as e:
|
||||
logger.warning(f"Error counting tokens: {e}")
|
||||
# 简单估算: 平均每个字符0.75个token
|
||||
return int(len(text) * 0.75)
|
||||
|
||||
async def count_messages(self, messages: list[dict[str, Any]], model: str | None = None) -> int:
|
||||
"""计算消息列表的Token数量"""
|
||||
if not self.enabled:
|
||||
return 0
|
||||
|
||||
model = model or self.default_model or "gpt-3.5-turbo"
|
||||
encoder = self._get_encoder(model)
|
||||
|
||||
# 获取每条消息的额外token数
|
||||
msg_overhead = self.MESSAGE_OVERHEAD.get(model, 3)
|
||||
|
||||
total_tokens = 0
|
||||
|
||||
for message in messages:
|
||||
# 每条消息的基本token
|
||||
total_tokens += msg_overhead
|
||||
|
||||
# 角色token
|
||||
role = message.get("role", "")
|
||||
if role:
|
||||
total_tokens += len(encoder.encode(role))
|
||||
|
||||
# 内容token
|
||||
content = message.get("content")
|
||||
if content:
|
||||
if isinstance(content, str):
|
||||
total_tokens += len(encoder.encode(content))
|
||||
elif isinstance(content, list):
|
||||
# 处理多模态内容
|
||||
for item in content:
|
||||
if item.get("type") == "text":
|
||||
text = item.get("text", "")
|
||||
total_tokens += len(encoder.encode(text))
|
||||
elif item.get("type") == "image_url":
|
||||
# 图像的token计算更复杂,这里简化处理
|
||||
# 低分辨率: 85 tokens, 高分辨率: 170 tokens
|
||||
detail = item.get("image_url", {}).get("detail", "auto")
|
||||
total_tokens += 170 if detail == "high" else 85
|
||||
|
||||
# 名称token
|
||||
name = message.get("name")
|
||||
if name:
|
||||
total_tokens += len(encoder.encode(name)) - 1 # name会减去1个token
|
||||
|
||||
# 工具调用
|
||||
tool_calls = message.get("tool_calls")
|
||||
if tool_calls:
|
||||
for tool_call in tool_calls:
|
||||
# 工具ID
|
||||
if "id" in tool_call:
|
||||
total_tokens += len(encoder.encode(tool_call["id"]))
|
||||
|
||||
# 函数信息
|
||||
function = tool_call.get("function", {})
|
||||
if "name" in function:
|
||||
total_tokens += len(encoder.encode(function["name"]))
|
||||
if "arguments" in function:
|
||||
total_tokens += len(encoder.encode(function["arguments"]))
|
||||
|
||||
# 添加固定的结束标记
|
||||
total_tokens += 3
|
||||
|
||||
return total_tokens
|
||||
|
||||
async def get_model_info(self, model: str) -> dict[str, Any]:
|
||||
"""获取模型信息"""
|
||||
info = {"model": model, "supported": self.supports_model(model)}
|
||||
|
||||
if self.supports_model(model):
|
||||
# 获取编码信息
|
||||
encoder = self._get_encoder(model)
|
||||
encoding_name = None
|
||||
|
||||
# 找到编码名称
|
||||
for m, enc in self.MODEL_ENCODINGS.items():
|
||||
if model.startswith(m):
|
||||
encoding_name = enc
|
||||
break
|
||||
|
||||
info.update(
|
||||
{
|
||||
"encoding": encoding_name or "unknown",
|
||||
"vocab_size": encoder.n_vocab if hasattr(encoder, "n_vocab") else None,
|
||||
"max_tokens": self._get_max_tokens(model),
|
||||
"message_overhead": self.MESSAGE_OVERHEAD.get(model, 3),
|
||||
}
|
||||
)
|
||||
|
||||
# 添加价格信息
|
||||
pricing = self.config.get("pricing", {})
|
||||
if model in pricing:
|
||||
info["pricing"] = pricing[model]
|
||||
|
||||
return info
|
||||
|
||||
def _get_max_tokens(self, model: str) -> int:
|
||||
"""获取模型的最大token数"""
|
||||
max_tokens_map = {
|
||||
"gpt-4": 8192,
|
||||
"gpt-4-32k": 32768,
|
||||
"gpt-4-turbo": 128000,
|
||||
"gpt-4o": 128000,
|
||||
"gpt-4o-mini": 128000,
|
||||
"gpt-3.5-turbo": 4096,
|
||||
"gpt-3.5-turbo-16k": 16384,
|
||||
"o1-preview": 128000,
|
||||
"o1-mini": 128000,
|
||||
}
|
||||
|
||||
# 完全匹配
|
||||
if model in max_tokens_map:
|
||||
return max_tokens_map[model]
|
||||
|
||||
# 前缀匹配
|
||||
for model_prefix, max_tokens in max_tokens_map.items():
|
||||
if model.startswith(model_prefix):
|
||||
return max_tokens
|
||||
|
||||
# 默认值
|
||||
return 4096
|
||||
|
||||
async def get_stats(self) -> dict[str, Any]:
|
||||
"""获取统计信息"""
|
||||
stats = await super().get_stats()
|
||||
stats.update(
|
||||
{
|
||||
"encoders_cached": _get_encoder_cached.cache_info().currsize,
|
||||
"tiktoken_available": TIKTOKEN_AVAILABLE,
|
||||
}
|
||||
)
|
||||
return stats
|
||||
Reference in New Issue
Block a user