mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +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:
@@ -97,7 +97,7 @@ class TestJWTTokenVerification:
|
||||
async def test_verify_expired_token_raises_error(self) -> None:
|
||||
"""测试验证过期令牌抛出异常"""
|
||||
# 创建一个已过期的 token
|
||||
data = {"sub": "user123", "type": "access"}
|
||||
data: dict[str, str | datetime] = {"sub": "user123", "type": "access"}
|
||||
expire = datetime.now(timezone.utc) - timedelta(hours=1)
|
||||
data["exp"] = expire
|
||||
expired_token = jwt.encode(data, JWT_SECRET_KEY, algorithm=JWT_ALGORITHM)
|
||||
@@ -269,6 +269,45 @@ class TestAPIKeyAuthentication:
|
||||
assert result[0] == mock_user
|
||||
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:
|
||||
"""测试 API Key 不存在"""
|
||||
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:
|
||||
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
|
||||
|
||||
@@ -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.billing_status == "settled"
|
||||
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