adai/MemoryBear
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge #12 into develop from feature/20251219_xjn

Browse Source 
feat(apikey system): api key authentication qps optimization

* feature/20251219_xjn: (4 commits)
  feat(apikey system): api key authentication adds the GET method
  feat(apikey system): api key authentication delete the GET method
  feat(apikey system): api key authentication qps optimization
  Merge branch 'refs/heads/develop' into feature/20251219_xjn

Signed-off-by: 谢俊男 <accounts_6853d0ea6f8174722fb0c8f1@mail.teambition.com>
Reviewed-by: zhuwenhui5566@163.com <zhuwenhui5566@163.com>
Merged-by: zhuwenhui5566@163.com <zhuwenhui5566@163.com>

CR-link: https://codeup.aliyun.com/redbearai/python/redbear-mem-open/change/12
This commit is contained in:
朱文辉
2025-12-18 17:19:53 +08:00
parent 64ecd8cabc f419b6b001
commit b15a73f3e7
2 changed files with 25 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

47
api/app/core/api_key_auth.py
View File

@@ -70,29 +70,6 @@ def require_api_key(
})
raise BusinessException("API Key 无效或已过期", BizCode.API_KEY_INVALID)
rate_limiter = RateLimiterService()
is_allowed, error_msg, rate_headers = await rate_limiter.check_all_limits(api_key_obj)
if not is_allowed:
logger.warning("API Key 限流触发", extra={
"api_key_id": str(api_key_obj.id),
"endpoint": str(request.url),
"method": request.method,
"error_msg": error_msg
})
# 根据错误消息判断限流类型
if "QPS" in error_msg:
code = BizCode.API_KEY_QPS_LIMIT_EXCEEDED
elif "Daily" in error_msg:
code = BizCode.API_KEY_DAILY_LIMIT_EXCEEDED
else:
code = BizCode.API_KEY_QUOTA_EXCEEDED
raise RateLimitException(
error_msg,
code,
rate_headers=rate_headers
)
if scopes:
missing_scopes = []
for scope in scopes:
@@ -138,6 +115,30 @@ def require_api_key(
scopes=api_key_obj.scopes,
resource_id=api_key_obj.resource_id,
)
rate_limiter = RateLimiterService()
is_allowed, error_msg, rate_headers = await rate_limiter.check_all_limits(api_key_obj)
if not is_allowed:
logger.warning("API Key 限流触发", extra={
"api_key_id": str(api_key_obj.id),
"endpoint": str(request.url),
"method": request.method,
"error_msg": error_msg
})
# 根据错误消息判断限流类型
if "QPS" in error_msg:
code = BizCode.API_KEY_QPS_LIMIT_EXCEEDED
elif "Daily" in error_msg:
code = BizCode.API_KEY_DAILY_LIMIT_EXCEEDED
else:
code = BizCode.API_KEY_QUOTA_EXCEEDED
raise RateLimitException(
error_msg,
code,
rate_headers=rate_headers
)
start_time = time.perf_counter()
response = await func(*args, **kwargs)
end_time = time.perf_counter()

2
api/app/services/api_key_service.py
View File

@@ -257,7 +257,7 @@ class RateLimiterService:
key = f"rate_limit:qps:{api_key_id}"
async with self.redis.pipeline() as pipe:
pipe.incr(key)
pipe.expire(key, 1) # 1 秒过期
pipe.expire(key, 1, nx=True) # 1 秒过期
results = await pipe.execute()
current = results[0]