mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +08:00
perf: cherry-pick PR #172 性能优化(不含 orjson)
- tiktoken 编码器改为 @lru_cache 全局缓存,避免多实例重复初始化 - 前缀匹配按长度排序,避免短前缀抢先匹配 - AuthService 更新 last_used_at 时临时关闭 expire_on_commit,减少额外 SELECT - UsageService.record_usage_batch 改用 bulk_insert_mappings 批量插入 - 已有记录查询增加 selectinload 预加载,避免 N+1 - Gemini normalizer 长行格式化 Co-Authored-By: AAEE86 <33052466+AAEE86@users.noreply.github.com>
This commit is contained in:
@@ -692,10 +692,18 @@ class GeminiNormalizer(FormatNormalizer):
|
|||||||
# 关闭前面的 thinking/text block(如果还开着)
|
# 关闭前面的 thinking/text block(如果还开着)
|
||||||
if ss.get("thinking_block_started") and not ss.get("thinking_block_stopped"):
|
if ss.get("thinking_block_started") and not ss.get("thinking_block_stopped"):
|
||||||
ss["thinking_block_stopped"] = True
|
ss["thinking_block_stopped"] = True
|
||||||
events.append(ContentBlockStopEvent(block_index=_reserve_block_index("thinking_block_index")))
|
events.append(
|
||||||
|
ContentBlockStopEvent(
|
||||||
|
block_index=_reserve_block_index("thinking_block_index")
|
||||||
|
)
|
||||||
|
)
|
||||||
if ss.get("text_block_started") and not ss.get("text_block_stopped"):
|
if ss.get("text_block_started") and not ss.get("text_block_stopped"):
|
||||||
ss["text_block_stopped"] = True
|
ss["text_block_stopped"] = True
|
||||||
events.append(ContentBlockStopEvent(block_index=_reserve_block_index("text_block_index")))
|
events.append(
|
||||||
|
ContentBlockStopEvent(
|
||||||
|
block_index=_reserve_block_index("text_block_index")
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
name = str(func_call.get("name") or "")
|
name = str(func_call.get("name") or "")
|
||||||
args = func_call.get("args")
|
args = func_call.get("args")
|
||||||
@@ -766,10 +774,14 @@ class GeminiNormalizer(FormatNormalizer):
|
|||||||
# 先补齐 content_block_stop(所有已开启的 block),再发送 MessageStop
|
# 先补齐 content_block_stop(所有已开启的 block),再发送 MessageStop
|
||||||
if ss.get("thinking_block_started") and not ss.get("thinking_block_stopped"):
|
if ss.get("thinking_block_started") and not ss.get("thinking_block_stopped"):
|
||||||
ss["thinking_block_stopped"] = True
|
ss["thinking_block_stopped"] = True
|
||||||
events.append(ContentBlockStopEvent(block_index=_reserve_block_index("thinking_block_index")))
|
events.append(
|
||||||
|
ContentBlockStopEvent(block_index=_reserve_block_index("thinking_block_index"))
|
||||||
|
)
|
||||||
if ss.get("text_block_started") and not ss.get("text_block_stopped"):
|
if ss.get("text_block_started") and not ss.get("text_block_stopped"):
|
||||||
ss["text_block_stopped"] = True
|
ss["text_block_stopped"] = True
|
||||||
events.append(ContentBlockStopEvent(block_index=_reserve_block_index("text_block_index")))
|
events.append(
|
||||||
|
ContentBlockStopEvent(block_index=_reserve_block_index("text_block_index"))
|
||||||
|
)
|
||||||
events.append(MessageStopEvent(stop_reason=stop_reason, usage=usage_info))
|
events.append(MessageStopEvent(stop_reason=stop_reason, usage=usage_info))
|
||||||
|
|
||||||
if "error" in chunk:
|
if "error" in chunk:
|
||||||
|
|||||||
@@ -5,6 +5,7 @@ Tiktoken Token计数插件
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from functools import lru_cache
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
from src.core.logger import logger
|
from src.core.logger import logger
|
||||||
@@ -16,11 +17,39 @@ try:
|
|||||||
import tiktoken
|
import tiktoken
|
||||||
|
|
||||||
TIKTOKEN_AVAILABLE = True
|
TIKTOKEN_AVAILABLE = True
|
||||||
except ImportError:
|
except ImportError: # pragma: no cover
|
||||||
TIKTOKEN_AVAILABLE = False
|
TIKTOKEN_AVAILABLE = False
|
||||||
tiktoken = None
|
tiktoken = None
|
||||||
|
|
||||||
|
|
||||||
|
@lru_cache(maxsize=256)
|
||||||
|
def _get_encoder_cached(model: str) -> Any:
|
||||||
|
"""全局编码器缓存。
|
||||||
|
|
||||||
|
目的:避免在多实例/多请求场景下重复初始化 tiktoken 编码器。
|
||||||
|
"""
|
||||||
|
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):
|
class TiktokenCounterPlugin(TokenCounterPlugin):
|
||||||
"""
|
"""
|
||||||
使用tiktoken库计算Token数量
|
使用tiktoken库计算Token数量
|
||||||
@@ -49,6 +78,13 @@ class TiktokenCounterPlugin(TokenCounterPlugin):
|
|||||||
"text-embedding-3-large": "cl100k_base",
|
"text-embedding-3-large": "cl100k_base",
|
||||||
}
|
}
|
||||||
|
|
||||||
|
# 前缀匹配顺序(从长到短)
|
||||||
|
MODEL_ENCODINGS_PREFIXES = sorted(
|
||||||
|
MODEL_ENCODINGS.items(),
|
||||||
|
key=lambda kv: len(kv[0]),
|
||||||
|
reverse=True,
|
||||||
|
)
|
||||||
|
|
||||||
# 每个消息的额外Token数
|
# 每个消息的额外Token数
|
||||||
MESSAGE_OVERHEAD = {
|
MESSAGE_OVERHEAD = {
|
||||||
"gpt-3.5-turbo": 4, # 每条消息
|
"gpt-3.5-turbo": 4, # 每条消息
|
||||||
@@ -84,42 +120,20 @@ class TiktokenCounterPlugin(TokenCounterPlugin):
|
|||||||
)
|
)
|
||||||
|
|
||||||
def _get_encoder(self, model: str) -> Any:
|
def _get_encoder(self, model: str) -> Any:
|
||||||
"""获取模型的编码器"""
|
"""获取模型的编码器(全局缓存)"""
|
||||||
if model in self._encoders:
|
return _get_encoder_cached(model)
|
||||||
return self._encoders[model]
|
|
||||||
|
|
||||||
# 获取编码名称
|
|
||||||
encoding_name = None
|
|
||||||
|
|
||||||
# 完全匹配
|
|
||||||
if model in self.MODEL_ENCODINGS:
|
|
||||||
encoding_name = self.MODEL_ENCODINGS[model]
|
|
||||||
else:
|
|
||||||
# 前缀匹配
|
|
||||||
for model_prefix, enc_name in self.MODEL_ENCODINGS.items():
|
|
||||||
if model.startswith(model_prefix):
|
|
||||||
encoding_name = enc_name
|
|
||||||
break
|
|
||||||
|
|
||||||
# 如果找不到,尝试使用模型名称
|
|
||||||
if not encoding_name:
|
|
||||||
try:
|
|
||||||
encoder = tiktoken.encoding_for_model(model)
|
|
||||||
self._encoders[model] = encoder
|
|
||||||
return encoder
|
|
||||||
except:
|
|
||||||
# 默认使用cl100k_base
|
|
||||||
encoding_name = "cl100k_base"
|
|
||||||
|
|
||||||
# 创建编码器
|
|
||||||
encoder = tiktoken.get_encoding(encoding_name)
|
|
||||||
self._encoders[model] = encoder
|
|
||||||
return encoder
|
|
||||||
|
|
||||||
def supports_model(self, model: str) -> bool:
|
def supports_model(self, model: str) -> bool:
|
||||||
"""检查是否支持指定模型"""
|
"""检查是否支持指定模型"""
|
||||||
# 支持所有OpenAI模型和一些兼容模型
|
# 支持所有OpenAI模型和一些兼容模型
|
||||||
openai_models = ["gpt-4", "gpt-3.5", "text-davinci", "text-embedding", "code-davinci", "o1"]
|
openai_models = [
|
||||||
|
"gpt-4",
|
||||||
|
"gpt-3.5",
|
||||||
|
"text-davinci",
|
||||||
|
"text-embedding",
|
||||||
|
"code-davinci",
|
||||||
|
"o1",
|
||||||
|
]
|
||||||
return any(model.startswith(prefix) for prefix in openai_models)
|
return any(model.startswith(prefix) for prefix in openai_models)
|
||||||
|
|
||||||
async def count_tokens(self, text: str, model: str | None = None) -> int:
|
async def count_tokens(self, text: str, model: str | None = None) -> int:
|
||||||
@@ -263,6 +277,9 @@ class TiktokenCounterPlugin(TokenCounterPlugin):
|
|||||||
"""获取统计信息"""
|
"""获取统计信息"""
|
||||||
stats = await super().get_stats()
|
stats = await super().get_stats()
|
||||||
stats.update(
|
stats.update(
|
||||||
{"encoders_cached": len(self._encoders), "tiktoken_available": TIKTOKEN_AVAILABLE}
|
{
|
||||||
|
"encoders_cached": _get_encoder_cached.cache_info().currsize,
|
||||||
|
"tiktoken_available": TIKTOKEN_AVAILABLE,
|
||||||
|
}
|
||||||
)
|
)
|
||||||
return stats
|
return stats
|
||||||
|
|||||||
@@ -477,7 +477,20 @@ class AuthService:
|
|||||||
# 更新最后使用时间(使用节流策略,减少数据库写入)
|
# 更新最后使用时间(使用节流策略,减少数据库写入)
|
||||||
if _should_update_last_used(key_record.id):
|
if _should_update_last_used(key_record.id):
|
||||||
key_record.last_used_at = datetime.now(timezone.utc)
|
key_record.last_used_at = datetime.now(timezone.utc)
|
||||||
db.commit() # 立即提交事务,释放数据库锁,避免阻塞后续请求
|
|
||||||
|
# 这里需要 commit 来尽快释放锁,但默认 expire_on_commit=True 会让已加载对象过期,
|
||||||
|
# 导致同一请求后续访问 user/api_key 字段时触发额外 SELECT。
|
||||||
|
original_expire_on_commit = getattr(db, "expire_on_commit", None)
|
||||||
|
try:
|
||||||
|
if original_expire_on_commit is not None:
|
||||||
|
db.expire_on_commit = False
|
||||||
|
db.commit() # 立即提交事务,释放数据库锁,避免阻塞后续请求
|
||||||
|
except Exception:
|
||||||
|
db.rollback()
|
||||||
|
raise
|
||||||
|
finally:
|
||||||
|
if original_expire_on_commit is not None:
|
||||||
|
db.expire_on_commit = original_expire_on_commit
|
||||||
|
|
||||||
api_key_fp = hashlib.sha256(api_key.encode()).hexdigest()[:12]
|
api_key_fp = hashlib.sha256(api_key.encode()).hexdigest()[:12]
|
||||||
logger.debug("API认证成功: 用户 {} (api_key_fp={})", user.email, api_key_fp)
|
logger.debug("API认证成功: 用户 {} (api_key_fp={})", user.email, api_key_fp)
|
||||||
|
|||||||
@@ -1450,7 +1450,17 @@ class UsageService:
|
|||||||
|
|
||||||
if request_ids:
|
if request_ids:
|
||||||
# 查询已存在的 Usage 记录(包括 pending/streaming 状态)
|
# 查询已存在的 Usage 记录(包括 pending/streaming 状态)
|
||||||
existing_records = db.query(Usage).filter(Usage.request_id.in_(request_ids)).all()
|
from sqlalchemy.orm import selectinload
|
||||||
|
|
||||||
|
existing_records = (
|
||||||
|
db.query(Usage)
|
||||||
|
.options(
|
||||||
|
selectinload(Usage.user),
|
||||||
|
selectinload(Usage.api_key),
|
||||||
|
)
|
||||||
|
.filter(Usage.request_id.in_(request_ids))
|
||||||
|
.all()
|
||||||
|
)
|
||||||
existing_usages = {u.request_id: u for u in existing_records}
|
existing_usages = {u.request_id: u for u in existing_records}
|
||||||
|
|
||||||
for record in records:
|
for record in records:
|
||||||
@@ -1643,7 +1653,10 @@ class UsageService:
|
|||||||
logger.warning("批量记录中更新失败: {}, request_id={}", e, request_id)
|
logger.warning("批量记录中更新失败: {}, request_id={}", e, request_id)
|
||||||
continue
|
continue
|
||||||
|
|
||||||
# 2. 处理需要新建的记录
|
# 2. 处理需要新建的记录(批量插入)
|
||||||
|
insert_mappings: list[dict[str, Any]] = []
|
||||||
|
insert_request_ids: list[str] = []
|
||||||
|
|
||||||
for i, (record, request_id, params) in enumerate(insert_params_list):
|
for i, (record, request_id, params) in enumerate(insert_params_list):
|
||||||
try:
|
try:
|
||||||
usage_params, total_cost, exc = insert_results[i]
|
usage_params, total_cost, exc = insert_results[i]
|
||||||
@@ -1653,16 +1666,17 @@ class UsageService:
|
|||||||
user = params.user
|
user = params.user
|
||||||
api_key = params.api_key
|
api_key = params.api_key
|
||||||
|
|
||||||
# 创建 Usage 记录
|
# 终态记录:补齐 settled/finalized_at;非终态:确保 billing_status=pending
|
||||||
usage = Usage(**usage_params)
|
status = usage_params.get("status")
|
||||||
# 新建记录默认 billing_status=settled,但补齐 finalized_at,便于审计与幂等判断
|
if status in terminal_statuses:
|
||||||
if usage_params.get("status") in terminal_statuses:
|
if usage_params.get("billing_status") in (None, "pending"):
|
||||||
if getattr(usage, "billing_status", None) in (None, "pending"):
|
usage_params["billing_status"] = "settled"
|
||||||
usage.billing_status = "settled"
|
usage_params.setdefault("finalized_at", finalized_at)
|
||||||
if getattr(usage, "finalized_at", None) is None:
|
elif usage_params.get("billing_status") is None:
|
||||||
usage.finalized_at = finalized_at
|
usage_params["billing_status"] = "pending"
|
||||||
db.add(usage)
|
|
||||||
usages.append(usage)
|
insert_mappings.append(usage_params)
|
||||||
|
insert_request_ids.append(request_id)
|
||||||
|
|
||||||
# 聚合统计
|
# 聚合统计
|
||||||
model_name = record.get("model") or "unknown"
|
model_name = record.get("model") or "unknown"
|
||||||
@@ -1689,6 +1703,24 @@ class UsageService:
|
|||||||
logger.warning("批量记录中跳过无效记录: {}, request_id={}", e, request_id)
|
logger.warning("批量记录中跳过无效记录: {}, request_id={}", e, request_id)
|
||||||
continue
|
continue
|
||||||
|
|
||||||
|
if insert_mappings:
|
||||||
|
try:
|
||||||
|
db.bulk_insert_mappings(Usage, insert_mappings)
|
||||||
|
|
||||||
|
# 仅用于保持返回值语义:将新建记录读回为 ORM 对象
|
||||||
|
inserted_records = (
|
||||||
|
db.query(Usage).filter(Usage.request_id.in_(insert_request_ids)).all()
|
||||||
|
)
|
||||||
|
inserted_map = {u.request_id: u for u in inserted_records}
|
||||||
|
for rid in insert_request_ids:
|
||||||
|
inserted_usage = inserted_map.get(rid)
|
||||||
|
if inserted_usage is not None:
|
||||||
|
usages.append(inserted_usage)
|
||||||
|
except Exception as e:
|
||||||
|
logger.error("批量插入 Usage 记录时出错: {}", e)
|
||||||
|
db.rollback()
|
||||||
|
raise
|
||||||
|
|
||||||
# 统计跳过的记录,失败率超过 10% 时提升日志级别
|
# 统计跳过的记录,失败率超过 10% 时提升日志级别
|
||||||
if skipped_count > 0:
|
if skipped_count > 0:
|
||||||
skip_ratio = skipped_count / total_count if total_count > 0 else 0
|
skip_ratio = skipped_count / total_count if total_count > 0 else 0
|
||||||
@@ -1763,13 +1795,14 @@ class UsageService:
|
|||||||
# 单次提交所有更改
|
# 单次提交所有更改
|
||||||
try:
|
try:
|
||||||
db.commit()
|
db.commit()
|
||||||
inserted_count = len(usages) - updated_count
|
inserted_count = len(insert_mappings)
|
||||||
|
total_written = updated_count + inserted_count
|
||||||
if updated_count > 0:
|
if updated_count > 0:
|
||||||
logger.debug(f"批量记录成功: 更新 {updated_count} 条, 新建 {inserted_count} 条")
|
logger.debug("批量记录成功: 更新 {} 条, 新建 {} 条", updated_count, inserted_count)
|
||||||
else:
|
else:
|
||||||
logger.debug(f"批量记录 {len(usages)} 条使用记录成功")
|
logger.debug("批量记录 {} 条使用记录成功", total_written)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"批量提交使用记录时出错: {e}")
|
logger.error("批量提交使用记录时出错: {}", e)
|
||||||
db.rollback()
|
db.rollback()
|
||||||
raise
|
raise
|
||||||
|
|
||||||
|
|||||||
@@ -97,7 +97,7 @@ class TestJWTTokenVerification:
|
|||||||
async def test_verify_expired_token_raises_error(self) -> None:
|
async def test_verify_expired_token_raises_error(self) -> None:
|
||||||
"""测试验证过期令牌抛出异常"""
|
"""测试验证过期令牌抛出异常"""
|
||||||
# 创建一个已过期的 token
|
# 创建一个已过期的 token
|
||||||
data = {"sub": "user123", "type": "access"}
|
data: dict[str, str | datetime] = {"sub": "user123", "type": "access"}
|
||||||
expire = datetime.now(timezone.utc) - timedelta(hours=1)
|
expire = datetime.now(timezone.utc) - timedelta(hours=1)
|
||||||
data["exp"] = expire
|
data["exp"] = expire
|
||||||
expired_token = jwt.encode(data, JWT_SECRET_KEY, algorithm=JWT_ALGORITHM)
|
expired_token = jwt.encode(data, JWT_SECRET_KEY, algorithm=JWT_ALGORITHM)
|
||||||
@@ -269,6 +269,45 @@ class TestAPIKeyAuthentication:
|
|||||||
assert result[0] == mock_user
|
assert result[0] == mock_user
|
||||||
assert result[1] == mock_api_key
|
assert result[1] == mock_api_key
|
||||||
|
|
||||||
|
def test_authenticate_api_key_last_used_commit_disables_expire_on_commit(self) -> None:
|
||||||
|
"""当需要更新 last_used_at 时,应临时关闭 expire_on_commit 以避免重复查询。"""
|
||||||
|
mock_user = MagicMock()
|
||||||
|
mock_user.id = "user-123"
|
||||||
|
mock_user.email = "test@example.com"
|
||||||
|
mock_user.is_active = True
|
||||||
|
mock_user.is_deleted = False
|
||||||
|
|
||||||
|
mock_api_key = MagicMock()
|
||||||
|
mock_api_key.id = "key-123"
|
||||||
|
mock_api_key.is_active = True
|
||||||
|
mock_api_key.is_locked = False
|
||||||
|
mock_api_key.expires_at = None
|
||||||
|
mock_api_key.user = mock_user
|
||||||
|
mock_api_key.balance_used_usd = 0.0
|
||||||
|
|
||||||
|
mock_db = MagicMock()
|
||||||
|
mock_db.expire_on_commit = True
|
||||||
|
|
||||||
|
def _commit_side_effect() -> None:
|
||||||
|
assert mock_db.expire_on_commit is False
|
||||||
|
|
||||||
|
mock_db.commit.side_effect = _commit_side_effect
|
||||||
|
mock_db.query.return_value.options.return_value.filter.return_value.first.return_value = (
|
||||||
|
mock_api_key
|
||||||
|
)
|
||||||
|
|
||||||
|
with patch("src.services.auth.service._should_update_last_used", return_value=True):
|
||||||
|
with patch("src.services.auth.service.ApiKey.hash_key", return_value="hashed_key"):
|
||||||
|
with patch(
|
||||||
|
"src.services.auth.service.ApiKeyService.check_balance",
|
||||||
|
return_value=(True, 100.0),
|
||||||
|
):
|
||||||
|
result = AuthService.authenticate_api_key(mock_db, "sk-test-key")
|
||||||
|
|
||||||
|
assert result is not None
|
||||||
|
assert mock_db.expire_on_commit is True
|
||||||
|
mock_db.commit.assert_called_once()
|
||||||
|
|
||||||
def test_authenticate_api_key_not_found(self) -> None:
|
def test_authenticate_api_key_not_found(self) -> None:
|
||||||
"""测试 API Key 不存在"""
|
"""测试 API Key 不存在"""
|
||||||
mock_db = MagicMock()
|
mock_db = MagicMock()
|
||||||
|
|||||||
@@ -1267,6 +1267,9 @@ async def test_record_usage_batch_updates_when_status_completed_billing_pending(
|
|||||||
def __init__(self, all_result: list[Any]) -> None:
|
def __init__(self, all_result: list[Any]) -> None:
|
||||||
self._all_result = all_result
|
self._all_result = all_result
|
||||||
|
|
||||||
|
def options(self, *args: Any, **kwargs: Any) -> "DummyQuery":
|
||||||
|
return self
|
||||||
|
|
||||||
def filter(self, *args: Any, **kwargs: Any) -> "DummyQuery":
|
def filter(self, *args: Any, **kwargs: Any) -> "DummyQuery":
|
||||||
return self
|
return self
|
||||||
|
|
||||||
@@ -1325,3 +1328,85 @@ async def test_record_usage_batch_updates_when_status_completed_billing_pending(
|
|||||||
assert existing.response_body == usage_params["response_body"]
|
assert existing.response_body == usage_params["response_body"]
|
||||||
assert existing.billing_status == "settled"
|
assert existing.billing_status == "settled"
|
||||||
assert existing.finalized_at is not None
|
assert existing.finalized_at is not None
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_record_usage_batch_uses_bulk_insert_mappings_for_new_records(
|
||||||
|
monkeypatch: Any,
|
||||||
|
) -> None:
|
||||||
|
"""确保批量新建 Usage 走 bulk_insert_mappings。"""
|
||||||
|
|
||||||
|
from src.models.database import Usage
|
||||||
|
from src.services.usage.service import UsageService
|
||||||
|
|
||||||
|
class DummyQuery:
|
||||||
|
def __init__(self, all_result: list[Any]) -> None:
|
||||||
|
self._all_result = all_result
|
||||||
|
|
||||||
|
def options(self, *args: Any, **kwargs: Any) -> "DummyQuery":
|
||||||
|
return self
|
||||||
|
|
||||||
|
def filter(self, *args: Any, **kwargs: Any) -> "DummyQuery":
|
||||||
|
return self
|
||||||
|
|
||||||
|
def all(self) -> list[Any]:
|
||||||
|
return self._all_result
|
||||||
|
|
||||||
|
inserted = Usage(
|
||||||
|
request_id="req-usage-batch-new",
|
||||||
|
provider_name="openai",
|
||||||
|
model="gpt-4",
|
||||||
|
status="completed",
|
||||||
|
billing_status="settled",
|
||||||
|
)
|
||||||
|
|
||||||
|
usage_query_calls = {"count": 0}
|
||||||
|
|
||||||
|
def _query_side_effect(model: Any) -> Any:
|
||||||
|
if model is Usage:
|
||||||
|
usage_query_calls["count"] += 1
|
||||||
|
if usage_query_calls["count"] == 1:
|
||||||
|
# existing_records
|
||||||
|
return DummyQuery([])
|
||||||
|
# inserted_records
|
||||||
|
return DummyQuery([inserted])
|
||||||
|
return DummyQuery([])
|
||||||
|
|
||||||
|
db = MagicMock()
|
||||||
|
db.query.side_effect = _query_side_effect
|
||||||
|
|
||||||
|
usage_params = {
|
||||||
|
"request_id": "req-usage-batch-new",
|
||||||
|
"provider_name": "openai",
|
||||||
|
"model": "gpt-4",
|
||||||
|
"status": "completed",
|
||||||
|
}
|
||||||
|
|
||||||
|
monkeypatch.setattr(
|
||||||
|
UsageService,
|
||||||
|
"_prepare_usage_records_batch",
|
||||||
|
AsyncMock(return_value=[(usage_params, 0.0, None)]),
|
||||||
|
)
|
||||||
|
|
||||||
|
result = await UsageService.record_usage_batch(
|
||||||
|
db,
|
||||||
|
[
|
||||||
|
{
|
||||||
|
"request_id": "req-usage-batch-new",
|
||||||
|
"provider": "openai",
|
||||||
|
"model": "gpt-4",
|
||||||
|
"status": "completed",
|
||||||
|
}
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|
||||||
|
db.bulk_insert_mappings.assert_called_once()
|
||||||
|
args, _kwargs = db.bulk_insert_mappings.call_args
|
||||||
|
assert args[0] is Usage
|
||||||
|
mappings = args[1]
|
||||||
|
assert isinstance(mappings, list) and len(mappings) == 1
|
||||||
|
assert mappings[0]["request_id"] == "req-usage-batch-new"
|
||||||
|
assert mappings[0]["billing_status"] == "settled"
|
||||||
|
assert mappings[0].get("finalized_at") is not None
|
||||||
|
|
||||||
|
assert result and result[0] is inserted
|
||||||
|
|||||||
41
tests/unit/test_tiktoken_counter_cache.py
Normal file
41
tests/unit/test_tiktoken_counter_cache.py
Normal file
@@ -0,0 +1,41 @@
|
|||||||
|
import pytest
|
||||||
|
|
||||||
|
|
||||||
|
class _DummyEncoder:
|
||||||
|
def encode(self, text: str) -> list[int]:
|
||||||
|
return [0] * len(text)
|
||||||
|
|
||||||
|
|
||||||
|
class _DummyTiktoken:
|
||||||
|
def __init__(self, calls: dict[str, int]) -> None:
|
||||||
|
self._calls = calls
|
||||||
|
|
||||||
|
def get_encoding(self, _name: str) -> _DummyEncoder:
|
||||||
|
self._calls["get_encoding"] += 1
|
||||||
|
return _DummyEncoder()
|
||||||
|
|
||||||
|
def encoding_for_model(self, _model: str) -> _DummyEncoder:
|
||||||
|
self._calls["encoding_for_model"] += 1
|
||||||
|
return _DummyEncoder()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_encoder_is_globally_cached(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||||
|
import src.plugins.token.tiktoken_counter as tc
|
||||||
|
|
||||||
|
# 清理缓存,避免受其他测试影响
|
||||||
|
tc._get_encoder_cached.cache_clear()
|
||||||
|
|
||||||
|
calls = {"get_encoding": 0, "encoding_for_model": 0}
|
||||||
|
|
||||||
|
monkeypatch.setattr(tc, "TIKTOKEN_AVAILABLE", True)
|
||||||
|
monkeypatch.setattr(tc, "tiktoken", _DummyTiktoken(calls))
|
||||||
|
|
||||||
|
p1 = tc.TiktokenCounterPlugin()
|
||||||
|
p2 = tc.TiktokenCounterPlugin()
|
||||||
|
|
||||||
|
# 两个实例对同一 model 请求编码器,底层 get_encoding 应只触发一次
|
||||||
|
await p1.count_tokens("hi", model="gpt-4")
|
||||||
|
await p2.count_tokens("hi", model="gpt-4")
|
||||||
|
|
||||||
|
assert calls["get_encoding"] == 1
|
||||||
Reference in New Issue
Block a user