mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +08:00
chore: 升级到 Python 3.14 并现代化代码
- 升级 Docker 基础镜像从 Python 3.12 到 3.14 - 更新 pyproject.toml 支持 Python 3.13/3.14 - 移除 Python 3.8/3.9/3.10/3.11 分类器 - 更新 black 和 mypy 配置目标版本 - 将 get_event_loop() 替换为 get_running_loop() 加上 RuntimeError 处理 - 简化 compute_cost_sync 中的 asyncio.run 使用 - Dict/List/Tuple/Set → dict/list/tuple/set (PEP 585) - Optional[T] → T | None (PEP 604) - Union[A, B] → A | B (PEP 604) - 移除废弃的 typing 导入 - 移除不必要的字符串引号注解
This commit is contained in:
@@ -6,8 +6,7 @@ JWT Token 黑名单服务
|
||||
|
||||
import hashlib
|
||||
import os
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Optional
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from src.clients.redis_client import get_redis_client
|
||||
from src.core.logger import logger
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
"""LDAP 认证服务"""
|
||||
|
||||
from typing import Any, Dict, Optional, Tuple
|
||||
from typing import Any
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
@@ -109,7 +109,7 @@ class LDAPService:
|
||||
"""LDAP 认证服务"""
|
||||
|
||||
@staticmethod
|
||||
def get_config(db: Session) -> Optional[LDAPConfig]:
|
||||
def get_config(db: Session) -> LDAPConfig | None:
|
||||
"""获取 LDAP 配置"""
|
||||
return db.query(LDAPConfig).first()
|
||||
|
||||
@@ -127,7 +127,7 @@ class LDAPService:
|
||||
return LDAPService.get_config_data(db) is not None
|
||||
|
||||
@staticmethod
|
||||
def get_config_data(db: Session) -> Optional[Dict[str, Any]]:
|
||||
def get_config_data(db: Session) -> dict[str, Any] | None:
|
||||
"""
|
||||
提前获取并解密配置,供线程池使用,避免跨线程共享 Session。
|
||||
|
||||
@@ -172,7 +172,7 @@ class LDAPService:
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def authenticate_with_config(config: Dict[str, Any], username: str, password: str) -> Optional[dict]:
|
||||
def authenticate_with_config(config: dict[str, Any], username: str, password: str) -> dict | None:
|
||||
"""
|
||||
LDAP bind 验证
|
||||
|
||||
@@ -306,7 +306,7 @@ class LDAPService:
|
||||
logger.warning(f"LDAP {name} 连接关闭失败: {e}")
|
||||
|
||||
@staticmethod
|
||||
def test_connection_with_config(config: Dict[str, Any]) -> Tuple[bool, str]:
|
||||
def test_connection_with_config(config: dict[str, Any]) -> tuple[bool, str]:
|
||||
"""
|
||||
测试 LDAP 连接
|
||||
|
||||
@@ -363,12 +363,12 @@ class LDAPService:
|
||||
|
||||
# 兼容旧接口:如果其他代码直接调用
|
||||
@staticmethod
|
||||
def authenticate(db: Session, username: str, password: str) -> Optional[dict]:
|
||||
def authenticate(db: Session, username: str, password: str) -> dict | None:
|
||||
config = LDAPService.get_config_data(db)
|
||||
return LDAPService.authenticate_with_config(config, username, password) if config else None
|
||||
|
||||
@staticmethod
|
||||
def test_connection(db: Session) -> Tuple[bool, str]:
|
||||
def test_connection(db: Session) -> tuple[bool, str]:
|
||||
config = LDAPService.get_config_data(db)
|
||||
if not config:
|
||||
return False, "LDAP 配置不存在或未启用"
|
||||
|
||||
@@ -1,7 +1,5 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
from typing import TYPE_CHECKING
|
||||
from urllib.parse import urlencode, urlparse, urlunparse
|
||||
|
||||
import httpx
|
||||
@@ -31,20 +29,20 @@ class OAuthProviderBase(ABC):
|
||||
userinfo_url: str
|
||||
default_scopes: tuple[str, ...] = ()
|
||||
|
||||
def get_effective_authorization_url(self, config: "OAuthProvider") -> str:
|
||||
def get_effective_authorization_url(self, config: OAuthProvider) -> str:
|
||||
return config.authorization_url_override or self.authorization_url
|
||||
|
||||
def get_effective_token_url(self, config: "OAuthProvider") -> str:
|
||||
def get_effective_token_url(self, config: OAuthProvider) -> str:
|
||||
return config.token_url_override or self.token_url
|
||||
|
||||
def get_effective_userinfo_url(self, config: "OAuthProvider") -> str:
|
||||
def get_effective_userinfo_url(self, config: OAuthProvider) -> str:
|
||||
return config.userinfo_url_override or self.userinfo_url
|
||||
|
||||
def get_effective_scopes(self, config: "OAuthProvider") -> str:
|
||||
def get_effective_scopes(self, config: OAuthProvider) -> str:
|
||||
scopes = config.scopes or list(self.default_scopes)
|
||||
return " ".join(scopes)
|
||||
|
||||
def get_authorization_url(self, config: "OAuthProvider", state: str) -> str:
|
||||
def get_authorization_url(self, config: OAuthProvider, state: str) -> str:
|
||||
"""
|
||||
构造 provider 授权 URL。
|
||||
|
||||
@@ -83,11 +81,11 @@ class OAuthProviderBase(ABC):
|
||||
return urlunparse(parsed._replace(query=urlencode(query)))
|
||||
|
||||
@abstractmethod
|
||||
async def exchange_code(self, config: "OAuthProvider", code: str) -> OAuthToken:
|
||||
async def exchange_code(self, config: OAuthProvider, code: str) -> OAuthToken:
|
||||
"""使用授权码兑换 token。"""
|
||||
|
||||
@abstractmethod
|
||||
async def get_user_info(self, config: "OAuthProvider", access_token: str) -> OAuthUserInfo:
|
||||
async def get_user_info(self, config: OAuthProvider, access_token: str) -> OAuthUserInfo:
|
||||
"""获取用户信息。"""
|
||||
|
||||
async def _http_post_form(
|
||||
@@ -96,7 +94,7 @@ class OAuthProviderBase(ABC):
|
||||
data: dict[str, str],
|
||||
*,
|
||||
timeout_seconds: float = 5.0,
|
||||
headers: Optional[dict[str, str]] = None,
|
||||
headers: dict[str, str] | None = None,
|
||||
) -> httpx.Response:
|
||||
async with httpx.AsyncClient(timeout=httpx.Timeout(timeout_seconds), verify=get_ssl_context()) as client:
|
||||
return await client.post(url, data=data, headers=headers)
|
||||
@@ -106,7 +104,7 @@ class OAuthProviderBase(ABC):
|
||||
url: str,
|
||||
*,
|
||||
timeout_seconds: float = 5.0,
|
||||
headers: Optional[dict[str, str]] = None,
|
||||
headers: dict[str, str] | None = None,
|
||||
) -> httpx.Response:
|
||||
async with httpx.AsyncClient(timeout=httpx.Timeout(timeout_seconds), verify=get_ssl_context()) as client:
|
||||
return await client.get(url, headers=headers)
|
||||
|
||||
@@ -1,27 +1,25 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Optional
|
||||
from typing import Any
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class OAuthToken:
|
||||
access_token: str
|
||||
token_type: str = "bearer"
|
||||
refresh_token: Optional[str] = None
|
||||
expires_in: Optional[int] = None
|
||||
id_token: Optional[str] = None
|
||||
scope: Optional[str] = None
|
||||
raw: Optional[dict[str, Any]] = None
|
||||
refresh_token: str | None = None
|
||||
expires_in: int | None = None
|
||||
id_token: str | None = None
|
||||
scope: str | None = None
|
||||
raw: dict[str, Any] | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class OAuthUserInfo:
|
||||
id: str
|
||||
username: Optional[str] = None
|
||||
email: Optional[str] = None
|
||||
email_verified: Optional[bool] = None
|
||||
raw: Optional[dict[str, Any]] = None
|
||||
username: str | None = None
|
||||
email: str | None = None
|
||||
email_verified: bool | None = None
|
||||
raw: dict[str, Any] | None = None
|
||||
|
||||
|
||||
class OAuthFlowError(Exception):
|
||||
|
||||
@@ -1,5 +1,3 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from src.core.logger import logger
|
||||
@@ -45,7 +43,7 @@ class LinuxDoOAuthProvider(OAuthProviderBase):
|
||||
# LinuxDo 不需要 scope
|
||||
default_scopes = ()
|
||||
|
||||
async def exchange_code(self, config: "OAuthProvider", code: str) -> OAuthToken:
|
||||
async def exchange_code(self, config: OAuthProvider, code: str) -> OAuthToken:
|
||||
url = self.get_effective_token_url(config)
|
||||
client_secret = config.get_client_secret()
|
||||
if not client_secret:
|
||||
@@ -86,7 +84,7 @@ class LinuxDoOAuthProvider(OAuthProviderBase):
|
||||
raw=data,
|
||||
)
|
||||
|
||||
async def get_user_info(self, config: "OAuthProvider", access_token: str) -> OAuthUserInfo:
|
||||
async def get_user_info(self, config: OAuthProvider, access_token: str) -> OAuthUserInfo:
|
||||
url = self.get_effective_userinfo_url(config)
|
||||
resp = await self._http_get(url, headers={"Authorization": f"Bearer {access_token}"})
|
||||
|
||||
|
||||
@@ -1,7 +1,4 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.services.auth.oauth.base import OAuthProviderBase
|
||||
@@ -22,7 +19,7 @@ class OAuthProviderRegistry:
|
||||
"""Provider 注册表(支持延迟 discover)。"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._providers: Dict[str, OAuthProviderBase] = {}
|
||||
self._providers: dict[str, OAuthProviderBase] = {}
|
||||
self._discovered: bool = False
|
||||
|
||||
def discover_providers(self) -> None:
|
||||
@@ -69,10 +66,10 @@ class OAuthProviderRegistry:
|
||||
def register(self, provider: OAuthProviderBase) -> None:
|
||||
self._providers[provider.provider_type] = provider
|
||||
|
||||
def get_provider(self, provider_type: str) -> Optional[OAuthProviderBase]:
|
||||
def get_provider(self, provider_type: str) -> OAuthProviderBase | None:
|
||||
return self._providers.get(provider_type)
|
||||
|
||||
def get_supported_types(self) -> List[SupportedOAuthType]:
|
||||
def get_supported_types(self) -> list[SupportedOAuthType]:
|
||||
return [
|
||||
SupportedOAuthType(
|
||||
provider_type=p.provider_type,
|
||||
@@ -86,7 +83,7 @@ class OAuthProviderRegistry:
|
||||
]
|
||||
|
||||
|
||||
_registry: Optional[OAuthProviderRegistry] = None
|
||||
_registry: OAuthProviderRegistry | None = None
|
||||
|
||||
|
||||
def get_oauth_provider_registry() -> OAuthProviderRegistry:
|
||||
|
||||
@@ -1,9 +1,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Optional
|
||||
from typing import Any
|
||||
from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse
|
||||
|
||||
import httpx
|
||||
@@ -39,7 +37,7 @@ class OAuthService:
|
||||
raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="OAuth 模块未启用")
|
||||
|
||||
@staticmethod
|
||||
def _get_provider_impl(provider_type: str) -> Optional[OAuthProviderBase]:
|
||||
def _get_provider_impl(provider_type: str) -> OAuthProviderBase | None:
|
||||
registry = get_oauth_provider_registry()
|
||||
registry.discover_providers()
|
||||
provider = registry.get_provider(provider_type)
|
||||
@@ -169,7 +167,7 @@ class OAuthService:
|
||||
return provider.get_authorization_url(config, state)
|
||||
|
||||
@staticmethod
|
||||
def _sanitize_username(raw: Optional[str]) -> str:
|
||||
def _sanitize_username(raw: str | None) -> str:
|
||||
if not raw or not raw.strip():
|
||||
return f"user_{uuid.uuid4().hex[:8]}"
|
||||
|
||||
@@ -224,7 +222,7 @@ class OAuthService:
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def _get_constraint_name(err: IntegrityError) -> Optional[str]:
|
||||
def _get_constraint_name(err: IntegrityError) -> str | None:
|
||||
orig = getattr(err, "orig", None)
|
||||
diag = getattr(orig, "diag", None)
|
||||
name = getattr(diag, "constraint_name", None)
|
||||
@@ -236,9 +234,9 @@ class OAuthService:
|
||||
db: Session,
|
||||
provider_type: str,
|
||||
state: str,
|
||||
code: Optional[str],
|
||||
error: Optional[str],
|
||||
error_description: Optional[str],
|
||||
code: str | None,
|
||||
error: str | None,
|
||||
error_description: str | None,
|
||||
) -> str:
|
||||
OAuthService._require_module_active(db)
|
||||
|
||||
@@ -405,8 +403,8 @@ class OAuthService:
|
||||
default_quota = SystemConfigService.get_config(db, "default_user_quota_usd", default=10.0)
|
||||
|
||||
# 生成唯一用户名 + 创建用户(简单重试)
|
||||
user: Optional[User] = None
|
||||
last_error: Optional[Exception] = None
|
||||
user: User | None = None
|
||||
last_error: Exception | None = None
|
||||
for _ in range(3):
|
||||
try:
|
||||
username = OAuthService._generate_unique_username(db, base_username)
|
||||
@@ -840,9 +838,9 @@ class OAuthService:
|
||||
async def test_provider_config_with_data(
|
||||
provider_type: str,
|
||||
client_id: str,
|
||||
client_secret: Optional[str],
|
||||
authorization_url_override: Optional[str],
|
||||
token_url_override: Optional[str],
|
||||
client_secret: str | None,
|
||||
authorization_url_override: str | None,
|
||||
token_url_override: str | None,
|
||||
redirect_uri: str,
|
||||
) -> dict[str, Any]:
|
||||
"""使用传入的表单数据测试配置,而非从数据库读取"""
|
||||
|
||||
@@ -4,7 +4,8 @@ import json
|
||||
import secrets
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Awaitable, Optional, cast
|
||||
from typing import Any, cast
|
||||
from collections.abc import Awaitable
|
||||
|
||||
from redis.asyncio import Redis
|
||||
|
||||
@@ -31,11 +32,11 @@ class OAuthStateData:
|
||||
nonce: str
|
||||
provider_type: str
|
||||
action: str # "login" | "bind"
|
||||
user_id: Optional[str]
|
||||
user_id: str | None
|
||||
created_at: int
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict[str, Any]) -> "OAuthStateData":
|
||||
def from_dict(cls, data: dict[str, Any]) -> OAuthStateData:
|
||||
return cls(
|
||||
nonce=str(data.get("nonce") or ""),
|
||||
provider_type=str(data.get("provider_type") or ""),
|
||||
@@ -50,7 +51,7 @@ def _state_key(nonce: str) -> str:
|
||||
|
||||
|
||||
async def create_oauth_state(
|
||||
redis: Redis, *, provider_type: str, action: str, user_id: Optional[str] = None
|
||||
redis: Redis, *, provider_type: str, action: str, user_id: str | None = None
|
||||
) -> str:
|
||||
nonce = secrets.token_urlsafe(24)
|
||||
data = {
|
||||
@@ -64,13 +65,13 @@ async def create_oauth_state(
|
||||
return nonce
|
||||
|
||||
|
||||
async def consume_oauth_state(redis: Redis, nonce: str) -> Optional[OAuthStateData]:
|
||||
async def consume_oauth_state(redis: Redis, nonce: str) -> OAuthStateData | None:
|
||||
if not nonce:
|
||||
return None
|
||||
|
||||
key = _state_key(nonce)
|
||||
# redis-py 的类型标注在 sync/async 之间会出现 Union;这里明确按 async 处理。
|
||||
raw = await cast(Awaitable[Optional[str]], redis.eval(CONSUME_STATE_SCRIPT, 1, key))
|
||||
raw = await cast(Awaitable[str | None], redis.eval(CONSUME_STATE_SCRIPT, 1, key))
|
||||
if not raw:
|
||||
return None
|
||||
|
||||
@@ -92,7 +93,7 @@ class OAuthBindTokenData:
|
||||
created_at: int
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict[str, Any]) -> "OAuthBindTokenData":
|
||||
def from_dict(cls, data: dict[str, Any]) -> OAuthBindTokenData:
|
||||
return cls(
|
||||
token=str(data.get("token") or ""),
|
||||
user_id=str(data.get("user_id") or ""),
|
||||
@@ -118,13 +119,13 @@ async def create_oauth_bind_token(redis: Redis, *, user_id: str, provider_type:
|
||||
return token
|
||||
|
||||
|
||||
async def consume_oauth_bind_token(redis: Redis, token: str) -> Optional[OAuthBindTokenData]:
|
||||
async def consume_oauth_bind_token(redis: Redis, token: str) -> OAuthBindTokenData | None:
|
||||
"""消费(验证并删除)OAuth 绑定令牌,返回令牌数据或 None"""
|
||||
if not token:
|
||||
return None
|
||||
|
||||
key = _bind_token_key(token)
|
||||
raw = await cast(Awaitable[Optional[str]], redis.eval(CONSUME_STATE_SCRIPT, 1, key))
|
||||
raw = await cast(Awaitable[str | None], redis.eval(CONSUME_STATE_SCRIPT, 1, key))
|
||||
if not raw:
|
||||
return None
|
||||
|
||||
|
||||
@@ -2,7 +2,6 @@
|
||||
认证服务
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import secrets
|
||||
@@ -11,7 +10,7 @@ import uuid
|
||||
from collections import OrderedDict
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from threading import Lock
|
||||
from typing import TYPE_CHECKING, Any, Dict, Optional
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import jwt
|
||||
from fastapi import HTTPException, status
|
||||
@@ -92,7 +91,7 @@ class AuthService:
|
||||
"""认证服务"""
|
||||
|
||||
@staticmethod
|
||||
def token_identity_matches_user(payload: Dict[str, Any], user: User) -> bool:
|
||||
def token_identity_matches_user(payload: dict[str, Any], user: User) -> bool:
|
||||
"""
|
||||
校验 token 的身份字段是否与用户一致。
|
||||
|
||||
@@ -147,7 +146,7 @@ class AuthService:
|
||||
return encoded_jwt
|
||||
|
||||
@staticmethod
|
||||
async def verify_token(token: str, token_type: Optional[str] = None) -> Dict[str, Any]:
|
||||
async def verify_token(token: str, token_type: str | None = None) -> dict[str, Any]:
|
||||
"""验证JWT令牌
|
||||
|
||||
Args:
|
||||
@@ -182,7 +181,7 @@ class AuthService:
|
||||
@staticmethod
|
||||
async def authenticate_user(
|
||||
db: Session, email: str, password: str, auth_type: str = "local"
|
||||
) -> Optional[User]:
|
||||
) -> User | None:
|
||||
"""用户登录认证
|
||||
|
||||
Args:
|
||||
@@ -219,7 +218,7 @@ class AuthService:
|
||||
),
|
||||
timeout=total_timeout,
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
except TimeoutError:
|
||||
logger.error(f"LDAP 认证总体超时({total_timeout}秒): {email}")
|
||||
return None
|
||||
|
||||
@@ -285,7 +284,7 @@ class AuthService:
|
||||
return user
|
||||
|
||||
@staticmethod
|
||||
async def _get_or_create_ldap_user(db: Session, ldap_user: dict) -> Optional[User]:
|
||||
async def _get_or_create_ldap_user(db: Session, ldap_user: dict) -> User | None:
|
||||
"""获取或创建 LDAP 用户
|
||||
|
||||
Args:
|
||||
@@ -299,7 +298,7 @@ class AuthService:
|
||||
|
||||
# 优先用稳定标识查找,避免邮箱变更/用户名冲突导致重复建号
|
||||
# 使用 with_for_update() 锁定行,防止并发创建
|
||||
user: Optional[User] = None
|
||||
user: User | None = None
|
||||
if ldap_dn:
|
||||
user = (
|
||||
db.query(User)
|
||||
@@ -416,7 +415,7 @@ class AuthService:
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def authenticate_api_key(db: Session, api_key: str) -> Optional[tuple[User, ApiKey]]:
|
||||
def authenticate_api_key(db: Session, api_key: str) -> tuple[User, ApiKey] | None:
|
||||
"""API密钥认证"""
|
||||
# 对API密钥进行哈希查找,预加载 user 关系以支持后续访问限制检查
|
||||
key_hash = ApiKey.hash_key(api_key)
|
||||
@@ -591,7 +590,7 @@ class AuthService:
|
||||
@staticmethod
|
||||
async def authenticate_management_token(
|
||||
db: Session, raw_token: str, client_ip: str
|
||||
) -> Optional[tuple[User, "ManagementToken"]]:
|
||||
) -> tuple[User, ManagementToken] | None:
|
||||
"""Management Token 认证
|
||||
|
||||
Args:
|
||||
|
||||
@@ -8,16 +8,15 @@
|
||||
- 自定义计费维度
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
from __future__ import annotations
|
||||
from typing import Any
|
||||
|
||||
from src.services.billing.models import (
|
||||
BillingDimension,
|
||||
BillingUnit,
|
||||
CostBreakdown,
|
||||
StandardizedUsage,
|
||||
)
|
||||
from src.services.billing.templates import (
|
||||
BILLING_TEMPLATE_REGISTRY,
|
||||
BillingTemplates,
|
||||
get_template,
|
||||
)
|
||||
@@ -50,8 +49,8 @@ class BillingCalculator:
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dimensions: Optional[List[BillingDimension]] = None,
|
||||
template: Optional[str] = None,
|
||||
dimensions: list[BillingDimension] | None = None,
|
||||
template: str | None = None,
|
||||
):
|
||||
"""
|
||||
初始化计费计算器
|
||||
@@ -73,10 +72,10 @@ class BillingCalculator:
|
||||
def calculate(
|
||||
self,
|
||||
usage: StandardizedUsage,
|
||||
prices: Dict[str, float],
|
||||
tiered_pricing: Optional[Dict[str, Any]] = None,
|
||||
cache_ttl_minutes: Optional[int] = None,
|
||||
total_input_context: Optional[int] = None,
|
||||
prices: dict[str, float],
|
||||
tiered_pricing: dict[str, Any] | None = None,
|
||||
cache_ttl_minutes: int | None = None,
|
||||
total_input_context: int | None = None,
|
||||
) -> CostBreakdown:
|
||||
"""
|
||||
计算费用
|
||||
@@ -131,9 +130,9 @@ class BillingCalculator:
|
||||
def _get_tier(
|
||||
self,
|
||||
usage: StandardizedUsage,
|
||||
tiered_pricing: Dict[str, Any],
|
||||
total_input_context: Optional[int] = None,
|
||||
) -> Tuple[Optional[Dict[str, Any]], Optional[int]]:
|
||||
tiered_pricing: dict[str, Any],
|
||||
total_input_context: int | None = None,
|
||||
) -> tuple[dict[str, Any] | None, int | None]:
|
||||
"""
|
||||
确定价格阶梯
|
||||
|
||||
@@ -178,9 +177,9 @@ class BillingCalculator:
|
||||
|
||||
def _get_cache_read_price_for_ttl(
|
||||
self,
|
||||
tier: Dict[str, Any],
|
||||
tier: dict[str, Any],
|
||||
cache_ttl_minutes: int,
|
||||
) -> Optional[float]:
|
||||
) -> float | None:
|
||||
"""
|
||||
根据缓存 TTL 获取缓存读取价格
|
||||
|
||||
@@ -212,7 +211,7 @@ class BillingCalculator:
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: Dict[str, Any]) -> "BillingCalculator":
|
||||
def from_config(cls, config: dict[str, Any]) -> BillingCalculator:
|
||||
"""
|
||||
从配置创建计费计算器
|
||||
|
||||
@@ -238,15 +237,15 @@ class BillingCalculator:
|
||||
|
||||
return cls(template=config.get("template", "claude"))
|
||||
|
||||
def get_dimension_names(self) -> List[str]:
|
||||
def get_dimension_names(self) -> list[str]:
|
||||
"""获取所有计费维度名称"""
|
||||
return [dim.name for dim in self.dimensions]
|
||||
|
||||
def get_required_price_fields(self) -> List[str]:
|
||||
def get_required_price_fields(self) -> list[str]:
|
||||
"""获取所需的价格字段名称"""
|
||||
return [dim.price_field for dim in self.dimensions]
|
||||
|
||||
def get_required_usage_fields(self) -> List[str]:
|
||||
def get_required_usage_fields(self) -> list[str]:
|
||||
"""获取所需的 usage 字段名称"""
|
||||
return [dim.usage_field for dim in self.dimensions]
|
||||
|
||||
@@ -258,14 +257,14 @@ def calculate_request_cost(
|
||||
cache_read_input_tokens: int,
|
||||
input_price_per_1m: float,
|
||||
output_price_per_1m: float,
|
||||
cache_creation_price_per_1m: Optional[float],
|
||||
cache_read_price_per_1m: Optional[float],
|
||||
price_per_request: Optional[float],
|
||||
tiered_pricing: Optional[Dict[str, Any]] = None,
|
||||
cache_ttl_minutes: Optional[int] = None,
|
||||
total_input_context: Optional[int] = None,
|
||||
cache_creation_price_per_1m: float | None,
|
||||
cache_read_price_per_1m: float | None,
|
||||
price_per_request: float | None,
|
||||
tiered_pricing: dict[str, Any] | None = None,
|
||||
cache_ttl_minutes: int | None = None,
|
||||
total_input_context: int | None = None,
|
||||
billing_template: str = "claude",
|
||||
) -> Dict[str, Any]:
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
计算请求成本的便捷函数
|
||||
|
||||
@@ -309,7 +308,7 @@ def calculate_request_cost(
|
||||
)
|
||||
|
||||
# 构建价格配置
|
||||
prices: Dict[str, float] = {
|
||||
prices: dict[str, float] = {
|
||||
"input_price_per_1m": input_price_per_1m,
|
||||
"output_price_per_1m": output_price_per_1m,
|
||||
}
|
||||
|
||||
@@ -8,9 +8,10 @@
|
||||
- CostBreakdown: 计费明细结果
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from typing import Any, Dict, Optional
|
||||
from typing import Any
|
||||
|
||||
|
||||
class BillingUnit(str, Enum):
|
||||
@@ -66,7 +67,7 @@ class BillingDimension:
|
||||
|
||||
return 0.0
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""转换为字典(用于序列化)"""
|
||||
return {
|
||||
"name": self.name,
|
||||
@@ -77,7 +78,7 @@ class BillingDimension:
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any]) -> "BillingDimension":
|
||||
def from_dict(cls, data: dict[str, Any]) -> BillingDimension:
|
||||
"""从字典创建实例"""
|
||||
return cls(
|
||||
name=data["name"],
|
||||
@@ -114,7 +115,7 @@ class StandardizedUsage:
|
||||
request_count: int = 1
|
||||
|
||||
# 扩展字段(未来可能需要的额外维度)
|
||||
extra: Dict[str, Any] = field(default_factory=dict)
|
||||
extra: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
def get(self, field_name: str, default: Any = 0) -> Any:
|
||||
"""
|
||||
@@ -149,9 +150,9 @@ class StandardizedUsage:
|
||||
else:
|
||||
self.extra[field_name] = value
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""转换为字典"""
|
||||
result: Dict[str, Any] = {
|
||||
result: dict[str, Any] = {
|
||||
"input_tokens": self.input_tokens,
|
||||
"output_tokens": self.output_tokens,
|
||||
"cache_creation_tokens": self.cache_creation_tokens,
|
||||
@@ -165,7 +166,7 @@ class StandardizedUsage:
|
||||
return result
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict[str, Any]) -> "StandardizedUsage":
|
||||
def from_dict(cls, data: dict[str, Any]) -> StandardizedUsage:
|
||||
"""从字典创建实例"""
|
||||
extra = data.pop("extra", {}) if "extra" in data else {}
|
||||
# 只取已知字段
|
||||
@@ -191,19 +192,19 @@ class CostBreakdown:
|
||||
"""
|
||||
|
||||
# 各维度费用 {"input": 0.01, "output": 0.02, "cache_read": 0.001, ...}
|
||||
costs: Dict[str, float] = field(default_factory=dict)
|
||||
costs: dict[str, float] = field(default_factory=dict)
|
||||
|
||||
# 总费用
|
||||
total_cost: float = 0.0
|
||||
|
||||
# 命中的阶梯索引(如果使用阶梯计费)
|
||||
tier_index: Optional[int] = None
|
||||
tier_index: int | None = None
|
||||
|
||||
# 货币单位
|
||||
currency: str = "USD"
|
||||
|
||||
# 使用的价格(用于记录和审计)
|
||||
effective_prices: Dict[str, float] = field(default_factory=dict)
|
||||
effective_prices: dict[str, float] = field(default_factory=dict)
|
||||
|
||||
# =========================================================================
|
||||
# 兼容旧接口的属性(便于渐进式迁移)
|
||||
@@ -244,7 +245,7 @@ class CostBreakdown:
|
||||
"""缓存存储费用(豆包等)"""
|
||||
return self.costs.get("cache_storage", 0.0)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""转换为字典"""
|
||||
return {
|
||||
"costs": self.costs,
|
||||
|
||||
@@ -9,7 +9,6 @@
|
||||
- PER_REQUEST: 按次计费
|
||||
"""
|
||||
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
from src.services.billing.models import BillingDimension, BillingUnit
|
||||
|
||||
@@ -25,7 +24,7 @@ class BillingTemplates:
|
||||
# - 缓存读取(约 0.1x 输入价格)
|
||||
# - 按次计费(可选,配置 price_per_request 时生效)
|
||||
# =========================================================================
|
||||
CLAUDE_STANDARD: List[BillingDimension] = [
|
||||
CLAUDE_STANDARD: list[BillingDimension] = [
|
||||
BillingDimension(
|
||||
name="input",
|
||||
usage_field="input_tokens",
|
||||
@@ -61,7 +60,7 @@ class BillingTemplates:
|
||||
# - 缓存读取(部分模型支持,无缓存创建费用)
|
||||
# - 按次计费(可选,配置 price_per_request 时生效)
|
||||
# =========================================================================
|
||||
OPENAI_STANDARD: List[BillingDimension] = [
|
||||
OPENAI_STANDARD: list[BillingDimension] = [
|
||||
BillingDimension(
|
||||
name="input",
|
||||
usage_field="input_tokens",
|
||||
@@ -95,7 +94,7 @@ class BillingTemplates:
|
||||
#
|
||||
# 注意:豆包的缓存创建是免费的,但存储需要按时付费
|
||||
# =========================================================================
|
||||
DOUBAO_STANDARD: List[BillingDimension] = [
|
||||
DOUBAO_STANDARD: list[BillingDimension] = [
|
||||
BillingDimension(
|
||||
name="input",
|
||||
usage_field="input_tokens",
|
||||
@@ -132,7 +131,7 @@ class BillingTemplates:
|
||||
# - 缓存读取
|
||||
# - 按次计费(用于图片生成等模型,需配置 price_per_request)
|
||||
# =========================================================================
|
||||
GEMINI_STANDARD: List[BillingDimension] = [
|
||||
GEMINI_STANDARD: list[BillingDimension] = [
|
||||
BillingDimension(
|
||||
name="input",
|
||||
usage_field="input_tokens",
|
||||
@@ -161,7 +160,7 @@ class BillingTemplates:
|
||||
# - 适用于某些图片生成模型、特殊 API 等
|
||||
# - 仅按请求次数计费,不按 token 计费
|
||||
# =========================================================================
|
||||
PER_REQUEST: List[BillingDimension] = [
|
||||
PER_REQUEST: list[BillingDimension] = [
|
||||
BillingDimension(
|
||||
name="request",
|
||||
usage_field="request_count",
|
||||
@@ -174,7 +173,7 @@ class BillingTemplates:
|
||||
# 混合计费(按次 + 按 token)
|
||||
# - 某些模型既有固定费用又有 token 费用
|
||||
# =========================================================================
|
||||
HYBRID_STANDARD: List[BillingDimension] = [
|
||||
HYBRID_STANDARD: list[BillingDimension] = [
|
||||
BillingDimension(
|
||||
name="input",
|
||||
usage_field="input_tokens",
|
||||
@@ -198,7 +197,7 @@ class BillingTemplates:
|
||||
# 模板注册表
|
||||
# =========================================================================
|
||||
|
||||
BILLING_TEMPLATE_REGISTRY: Dict[str, List[BillingDimension]] = {
|
||||
BILLING_TEMPLATE_REGISTRY: dict[str, list[BillingDimension]] = {
|
||||
# 按厂商名称
|
||||
"claude": BillingTemplates.CLAUDE_STANDARD,
|
||||
"anthropic": BillingTemplates.CLAUDE_STANDARD,
|
||||
@@ -215,7 +214,7 @@ BILLING_TEMPLATE_REGISTRY: Dict[str, List[BillingDimension]] = {
|
||||
}
|
||||
|
||||
|
||||
def get_template(name: Optional[str]) -> List[BillingDimension]:
|
||||
def get_template(name: str | None) -> list[BillingDimension]:
|
||||
"""
|
||||
获取计费模板
|
||||
|
||||
@@ -236,6 +235,6 @@ def get_template(name: Optional[str]) -> List[BillingDimension]:
|
||||
return template
|
||||
|
||||
|
||||
def list_templates() -> List[str]:
|
||||
def list_templates() -> list[str]:
|
||||
"""列出所有可用的模板名称"""
|
||||
return list(BILLING_TEMPLATE_REGISTRY.keys())
|
||||
|
||||
@@ -9,7 +9,7 @@ Usage 字段映射器
|
||||
- GEMINI / GEMINI_CLI: Google Gemini API
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, Optional
|
||||
from typing import Any
|
||||
|
||||
from src.services.billing.models import StandardizedUsage
|
||||
|
||||
@@ -47,7 +47,7 @@ class UsageMapper:
|
||||
# =========================================================================
|
||||
|
||||
# OpenAI 格式字段映射
|
||||
OPENAI_MAPPING: Dict[str, str] = {
|
||||
OPENAI_MAPPING: dict[str, str] = {
|
||||
"prompt_tokens": "input_tokens",
|
||||
"completion_tokens": "output_tokens",
|
||||
"prompt_tokens_details.cached_tokens": "cache_read_tokens",
|
||||
@@ -55,7 +55,7 @@ class UsageMapper:
|
||||
}
|
||||
|
||||
# Claude 格式字段映射
|
||||
CLAUDE_MAPPING: Dict[str, str] = {
|
||||
CLAUDE_MAPPING: dict[str, str] = {
|
||||
"input_tokens": "input_tokens",
|
||||
"output_tokens": "output_tokens",
|
||||
"cache_creation_input_tokens": "cache_creation_tokens",
|
||||
@@ -63,7 +63,7 @@ class UsageMapper:
|
||||
}
|
||||
|
||||
# Gemini 格式字段映射
|
||||
GEMINI_MAPPING: Dict[str, str] = {
|
||||
GEMINI_MAPPING: dict[str, str] = {
|
||||
"promptTokenCount": "input_tokens",
|
||||
"candidatesTokenCount": "output_tokens",
|
||||
"cachedContentTokenCount": "cache_read_tokens",
|
||||
@@ -74,7 +74,7 @@ class UsageMapper:
|
||||
}
|
||||
|
||||
# 格式名称到映射的对应关系
|
||||
FORMAT_MAPPINGS: Dict[str, Dict[str, str]] = {
|
||||
FORMAT_MAPPINGS: dict[str, dict[str, str]] = {
|
||||
"OPENAI": OPENAI_MAPPING,
|
||||
"OPENAI_CLI": OPENAI_MAPPING,
|
||||
"CLAUDE": CLAUDE_MAPPING,
|
||||
@@ -86,9 +86,9 @@ class UsageMapper:
|
||||
@classmethod
|
||||
def map(
|
||||
cls,
|
||||
raw_usage: Dict[str, Any],
|
||||
raw_usage: dict[str, Any],
|
||||
api_format: str,
|
||||
extra_mapping: Optional[Dict[str, str]] = None,
|
||||
extra_mapping: dict[str, str] | None = None,
|
||||
) -> StandardizedUsage:
|
||||
"""
|
||||
将原始 usage 映射为标准化格式
|
||||
@@ -124,7 +124,7 @@ class UsageMapper:
|
||||
@classmethod
|
||||
def map_from_response(
|
||||
cls,
|
||||
response: Dict[str, Any],
|
||||
response: dict[str, Any],
|
||||
api_format: str,
|
||||
) -> StandardizedUsage:
|
||||
"""
|
||||
@@ -145,7 +145,7 @@ class UsageMapper:
|
||||
format_upper = api_format.upper() if api_format else ""
|
||||
|
||||
# 提取 usage 部分
|
||||
usage_data: Dict[str, Any] = {}
|
||||
usage_data: dict[str, Any] = {}
|
||||
|
||||
if format_upper.startswith("GEMINI"):
|
||||
# Gemini: usageMetadata
|
||||
@@ -162,7 +162,7 @@ class UsageMapper:
|
||||
return cls.map(usage_data, api_format)
|
||||
|
||||
@classmethod
|
||||
def _get_mapping(cls, api_format: str) -> Dict[str, str]:
|
||||
def _get_mapping(cls, api_format: str) -> dict[str, str]:
|
||||
"""获取对应格式的字段映射"""
|
||||
if not api_format:
|
||||
return cls.CLAUDE_MAPPING
|
||||
@@ -182,7 +182,7 @@ class UsageMapper:
|
||||
return cls.CLAUDE_MAPPING
|
||||
|
||||
@classmethod
|
||||
def _get_nested_value(cls, data: Dict[str, Any], path: str) -> Any:
|
||||
def _get_nested_value(cls, data: dict[str, Any], path: str) -> Any:
|
||||
"""
|
||||
获取嵌套字段值
|
||||
|
||||
@@ -212,7 +212,7 @@ class UsageMapper:
|
||||
return value
|
||||
|
||||
@classmethod
|
||||
def register_format(cls, format_name: str, mapping: Dict[str, str]) -> None:
|
||||
def register_format(cls, format_name: str, mapping: dict[str, str]) -> None:
|
||||
"""
|
||||
注册新的格式映射
|
||||
|
||||
@@ -234,7 +234,7 @@ class UsageMapper:
|
||||
|
||||
|
||||
def map_usage(
|
||||
raw_usage: Dict[str, Any],
|
||||
raw_usage: dict[str, Any],
|
||||
api_format: str,
|
||||
) -> StandardizedUsage:
|
||||
"""
|
||||
@@ -251,7 +251,7 @@ def map_usage(
|
||||
|
||||
|
||||
def map_usage_from_response(
|
||||
response: Dict[str, Any],
|
||||
response: dict[str, Any],
|
||||
api_format: str,
|
||||
) -> StandardizedUsage:
|
||||
"""
|
||||
|
||||
42
src/services/cache/affinity_manager.py
vendored
42
src/services/cache/affinity_manager.py
vendored
@@ -23,7 +23,7 @@ import json
|
||||
import os
|
||||
import time
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Any, Dict, List, NamedTuple, Optional, Tuple
|
||||
from typing import Any, NamedTuple
|
||||
|
||||
from src.config.constants import CacheTTL
|
||||
from src.core.logger import logger
|
||||
@@ -86,18 +86,18 @@ class CacheAffinityManager:
|
||||
"""
|
||||
self.redis = redis_client
|
||||
self.default_ttl = default_ttl
|
||||
self._memory_store: Dict[str, Dict[str, Any]] = {}
|
||||
self._memory_lock: Optional[asyncio.Lock] = None
|
||||
self._memory_store: dict[str, dict[str, Any]] = {}
|
||||
self._memory_lock: asyncio.Lock | None = None
|
||||
|
||||
# L1 缓存(即使使用 Redis 也启用,减少网络往返)
|
||||
self._l1_cache_ttl = int(os.getenv("CACHE_AFFINITY_L1_TTL", str(CacheTTL.L1_LOCAL)))
|
||||
self._l1_cache: Dict[str, Tuple[float, Dict[str, Any]]] = {}
|
||||
self._l1_cache: dict[str, tuple[float, dict[str, Any]]] = {}
|
||||
self._l1_lock = asyncio.Lock()
|
||||
self._l1_max_size = int(os.getenv("CACHE_AFFINITY_L1_MAX_SIZE", "1000")) # 最大缓存条目数
|
||||
self._l1_last_cleanup = time.time()
|
||||
|
||||
# 请求级别锁,避免同一用户+端点同时更新造成抖动
|
||||
self._request_locks: Dict[str, asyncio.Lock] = {}
|
||||
self._request_locks: dict[str, asyncio.Lock] = {}
|
||||
|
||||
# 统计信息
|
||||
self._stats = {
|
||||
@@ -138,7 +138,7 @@ class CacheAffinityManager:
|
||||
"""
|
||||
return f"cache_affinity:{affinity_key}:{api_format}:{model_name}"
|
||||
|
||||
async def _get_l1_entry(self, cache_key: str) -> Optional[Dict[str, Any]]:
|
||||
async def _get_l1_entry(self, cache_key: str) -> dict[str, Any] | None:
|
||||
async with self._l1_lock:
|
||||
record = self._l1_cache.get(cache_key)
|
||||
if not record:
|
||||
@@ -149,7 +149,7 @@ class CacheAffinityManager:
|
||||
return None
|
||||
return dict(payload)
|
||||
|
||||
async def _set_l1_entry(self, cache_key: str, payload: Optional[Dict[str, Any]]):
|
||||
async def _set_l1_entry(self, cache_key: str, payload: dict[str, Any] | None):
|
||||
async with self._l1_lock:
|
||||
if not payload:
|
||||
self._l1_cache.pop(cache_key, None)
|
||||
@@ -205,7 +205,7 @@ class CacheAffinityManager:
|
||||
finally:
|
||||
lock.release()
|
||||
|
||||
async def _load_affinity_dict(self, cache_key: str) -> Optional[Dict[str, Any]]:
|
||||
async def _load_affinity_dict(self, cache_key: str) -> dict[str, Any] | None:
|
||||
"""读取缓存亲和性字典"""
|
||||
# 先尝试L1缓存
|
||||
l1_value = await self._get_l1_entry(cache_key)
|
||||
@@ -228,7 +228,7 @@ class CacheAffinityManager:
|
||||
return dict(record) if record else None
|
||||
|
||||
async def _save_affinity_dict(
|
||||
self, cache_key: str, ttl: int, affinity_dict: Dict[str, Any]
|
||||
self, cache_key: str, ttl: int, affinity_dict: dict[str, Any]
|
||||
) -> None:
|
||||
"""存储缓存亲和性字典"""
|
||||
if not self._is_memory_backend():
|
||||
@@ -252,7 +252,7 @@ class CacheAffinityManager:
|
||||
|
||||
await self._set_l1_entry(cache_key, None)
|
||||
|
||||
async def _snapshot_memory_items(self) -> Dict[str, Dict[str, Any]]:
|
||||
async def _snapshot_memory_items(self) -> dict[str, dict[str, Any]]:
|
||||
"""复制内存存储内容(仅内存模式使用)"""
|
||||
lock = self._get_memory_lock()
|
||||
async with lock:
|
||||
@@ -260,7 +260,7 @@ class CacheAffinityManager:
|
||||
|
||||
async def get_affinity(
|
||||
self, affinity_key: str, api_format: str, model_name: str
|
||||
) -> Optional[CacheAffinity]:
|
||||
) -> CacheAffinity | None:
|
||||
"""
|
||||
获取指定亲和性标识符对特定API格式和模型的缓存亲和性
|
||||
|
||||
@@ -315,7 +315,7 @@ class CacheAffinityManager:
|
||||
api_format: str,
|
||||
model_name: str,
|
||||
supports_caching: bool = True,
|
||||
ttl: Optional[int] = None,
|
||||
ttl: int | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
设置指定亲和性标识符对特定API格式和模型的缓存亲和性
|
||||
@@ -345,7 +345,7 @@ class CacheAffinityManager:
|
||||
try:
|
||||
async with self._acquire_request_lock(cache_key):
|
||||
existing_dict = await self._load_affinity_dict(cache_key)
|
||||
existing_affinity: Optional[CacheAffinity] = None
|
||||
existing_affinity: CacheAffinity | None = None
|
||||
if existing_dict and current_time <= existing_dict.get("expire_at", 0):
|
||||
existing_affinity = CacheAffinity(
|
||||
provider_id=existing_dict["provider_id"],
|
||||
@@ -408,9 +408,9 @@ class CacheAffinityManager:
|
||||
affinity_key: str,
|
||||
api_format: str,
|
||||
model_name: str,
|
||||
key_id: Optional[str] = None,
|
||||
provider_id: Optional[str] = None,
|
||||
endpoint_id: Optional[str] = None,
|
||||
key_id: str | None = None,
|
||||
provider_id: str | None = None,
|
||||
endpoint_id: str | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
失效指定亲和性标识符对特定API格式和模型的缓存亲和性
|
||||
@@ -527,7 +527,7 @@ class CacheAffinityManager:
|
||||
logger.exception(f"清除缓存亲和性失败: {e}")
|
||||
return 0
|
||||
|
||||
def get_stats(self) -> Dict[str, Any]:
|
||||
def get_stats(self) -> dict[str, Any]:
|
||||
"""获取统计信息"""
|
||||
cache_hit_rate = 0.0
|
||||
total_requests = self._stats["cache_hits"] + self._stats["cache_misses"]
|
||||
@@ -551,7 +551,7 @@ class CacheAffinityManager:
|
||||
},
|
||||
}
|
||||
|
||||
async def list_affinities(self) -> List[Dict[str, Any]]:
|
||||
async def list_affinities(self) -> list[dict[str, Any]]:
|
||||
"""获取所有缓存亲和性列表
|
||||
|
||||
返回的每条记录包含:
|
||||
@@ -560,7 +560,7 @@ class CacheAffinityManager:
|
||||
- api_format, model_name: API 格式和模型名称
|
||||
- created_at, expire_at, request_count: 缓存元数据
|
||||
"""
|
||||
results: List[Dict[str, Any]] = []
|
||||
results: list[dict[str, Any]] = []
|
||||
|
||||
try:
|
||||
pattern = "cache_affinity:*"
|
||||
@@ -605,7 +605,7 @@ class CacheAffinityManager:
|
||||
break
|
||||
else:
|
||||
snapshot = await self._snapshot_memory_items()
|
||||
expired_keys: List[str] = []
|
||||
expired_keys: list[str] = []
|
||||
current_time = time.time()
|
||||
|
||||
for cache_key, affinity in snapshot.items():
|
||||
@@ -644,7 +644,7 @@ class CacheAffinityManager:
|
||||
|
||||
|
||||
# 全局单例
|
||||
_affinity_manager: Optional[CacheAffinityManager] = None
|
||||
_affinity_manager: CacheAffinityManager | None = None
|
||||
|
||||
|
||||
async def get_affinity_manager(redis_client=None) -> CacheAffinityManager:
|
||||
|
||||
176
src/services/cache/aware_scheduler.py
vendored
176
src/services/cache/aware_scheduler.py
vendored
@@ -28,14 +28,13 @@
|
||||
- 失效缓存亲和性,避免重复选择故障资源
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import random
|
||||
import re
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple, Union
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from sqlalchemy.orm import Session, selectinload
|
||||
|
||||
@@ -65,7 +64,6 @@ from src.services.rate_limit.adaptive_reservation import (
|
||||
get_adaptive_reservation_manager,
|
||||
)
|
||||
from src.services.rate_limit.concurrency_manager import get_concurrency_manager
|
||||
from src.services.system.config import SystemConfigService
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -77,8 +75,8 @@ class ProviderCandidate:
|
||||
key: ProviderAPIKey
|
||||
is_cached: bool = False
|
||||
is_skipped: bool = False # 是否被跳过
|
||||
skip_reason: Optional[str] = None # 跳过原因
|
||||
mapping_matched_model: Optional[str] = None # 通过映射匹配到的模型名(用于实际请求)
|
||||
skip_reason: str | None = None # 跳过原因
|
||||
mapping_matched_model: str | None = None # 通过映射匹配到的模型名(用于实际请求)
|
||||
needs_conversion: bool = False # 是否需要格式转换
|
||||
provider_api_format: str = "" # Provider 端点实际格式(用于健康度/熔断 bucket)
|
||||
|
||||
@@ -86,7 +84,7 @@ class ProviderCandidate:
|
||||
@dataclass
|
||||
class ConcurrencySnapshot:
|
||||
key_current: int
|
||||
key_limit: Optional[int]
|
||||
key_limit: int | None
|
||||
is_cached_user: bool = False
|
||||
# 动态预留信息
|
||||
reservation_ratio: float = 0.0
|
||||
@@ -134,8 +132,8 @@ class CacheAwareScheduler:
|
||||
def __init__(
|
||||
self,
|
||||
redis_client=None,
|
||||
priority_mode: Optional[str] = None,
|
||||
scheduling_mode: Optional[str] = None,
|
||||
priority_mode: str | None = None,
|
||||
scheduling_mode: str | None = None,
|
||||
):
|
||||
"""
|
||||
初始化调度器
|
||||
@@ -160,7 +158,7 @@ class CacheAwareScheduler:
|
||||
)
|
||||
|
||||
# 初始化子组件(将在第一次使用时异步初始化)
|
||||
self._affinity_manager: Optional[CacheAffinityManager] = None
|
||||
self._affinity_manager: CacheAffinityManager | None = None
|
||||
self._concurrency_manager = None
|
||||
# 动态预留管理器(同步初始化)
|
||||
self._reservation_manager: AdaptiveReservationManager = get_adaptive_reservation_manager()
|
||||
@@ -194,13 +192,13 @@ class CacheAwareScheduler:
|
||||
self,
|
||||
db: Session,
|
||||
affinity_key: str,
|
||||
api_format: Union[str, APIFormat],
|
||||
api_format: str | APIFormat,
|
||||
model_name: str,
|
||||
excluded_endpoints: Optional[List[str]] = None,
|
||||
excluded_keys: Optional[List[str]] = None,
|
||||
excluded_endpoints: list[str] | None = None,
|
||||
excluded_keys: list[str] | None = None,
|
||||
provider_batch_size: int = 20,
|
||||
max_candidates_per_batch: Optional[int] = None,
|
||||
) -> Tuple[Provider, ProviderEndpoint, ProviderAPIKey]:
|
||||
max_candidates_per_batch: int | None = None,
|
||||
) -> tuple[Provider, ProviderEndpoint, ProviderAPIKey]:
|
||||
"""
|
||||
缓存感知选择 - 核心方法
|
||||
|
||||
@@ -318,7 +316,7 @@ class CacheAwareScheduler:
|
||||
|
||||
raise ProviderNotAvailableException("服务暂时繁忙,请稍后重试")
|
||||
|
||||
def _get_effective_rpm_limit(self, key: ProviderAPIKey) -> Optional[int]:
|
||||
def _get_effective_rpm_limit(self, key: ProviderAPIKey) -> int | None:
|
||||
"""
|
||||
获取有效的 RPM 限制
|
||||
|
||||
@@ -348,7 +346,7 @@ class CacheAwareScheduler:
|
||||
self,
|
||||
key: ProviderAPIKey,
|
||||
is_cached_user: bool = False,
|
||||
) -> Tuple[bool, ConcurrencySnapshot]:
|
||||
) -> tuple[bool, ConcurrencySnapshot]:
|
||||
"""
|
||||
检查 RPM 限制是否可用(使用动态预留机制)
|
||||
|
||||
@@ -448,7 +446,7 @@ class CacheAwareScheduler:
|
||||
)
|
||||
can_use = False
|
||||
|
||||
key_limit_for_snapshot: Optional[int]
|
||||
key_limit_for_snapshot: int | None
|
||||
if is_cached_user:
|
||||
key_limit_for_snapshot = effective_key_limit
|
||||
elif effective_key_limit is not None:
|
||||
@@ -471,8 +469,8 @@ class CacheAwareScheduler:
|
||||
|
||||
def _get_effective_restrictions(
|
||||
self,
|
||||
user_api_key: Optional[ApiKey],
|
||||
) -> Dict[str, Any]:
|
||||
user_api_key: ApiKey | None,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
获取有效的访问限制(合并 ApiKey 和 User 的限制)
|
||||
|
||||
@@ -549,16 +547,16 @@ class CacheAwareScheduler:
|
||||
async def list_all_candidates(
|
||||
self,
|
||||
db: Session,
|
||||
api_format: Union[str, APIFormat],
|
||||
api_format: str | APIFormat,
|
||||
model_name: str,
|
||||
affinity_key: Optional[str] = None,
|
||||
user_api_key: Optional[ApiKey] = None,
|
||||
affinity_key: str | None = None,
|
||||
user_api_key: ApiKey | None = None,
|
||||
provider_offset: int = 0,
|
||||
provider_limit: Optional[int] = None,
|
||||
max_candidates: Optional[int] = None,
|
||||
provider_limit: int | None = None,
|
||||
max_candidates: int | None = None,
|
||||
is_stream: bool = False,
|
||||
capability_requirements: Optional[Dict[str, bool]] = None,
|
||||
) -> Tuple[List[ProviderCandidate], str]:
|
||||
capability_requirements: dict[str, bool] | None = None,
|
||||
) -> tuple[list[ProviderCandidate], str]:
|
||||
"""
|
||||
预先获取所有可用的 Provider/Endpoint/Key 组合
|
||||
|
||||
@@ -599,7 +597,7 @@ class CacheAwareScheduler:
|
||||
global_model_id: str = str(global_model.id)
|
||||
|
||||
# 提取模型映射(用于 Provider Key 的 allowed_models 匹配)
|
||||
model_mappings: List[str] = (global_model.config or {}).get("model_mappings", [])
|
||||
model_mappings: list[str] = (global_model.config or {}).get("model_mappings", [])
|
||||
if model_mappings:
|
||||
logger.debug(
|
||||
f"[Scheduler] GlobalModel={global_model.name} 配置了映射规则: {model_mappings}"
|
||||
@@ -712,8 +710,8 @@ class CacheAwareScheduler:
|
||||
self,
|
||||
db: Session,
|
||||
provider_offset: int = 0,
|
||||
provider_limit: Optional[int] = None,
|
||||
) -> List[Provider]:
|
||||
provider_limit: int | None = None,
|
||||
) -> list[Provider]:
|
||||
"""
|
||||
查询活跃的 Providers(带预加载)
|
||||
|
||||
@@ -751,10 +749,10 @@ class CacheAwareScheduler:
|
||||
db: Session,
|
||||
provider: Provider,
|
||||
model_name: str,
|
||||
api_format: Optional[str] = None,
|
||||
api_format: str | None = None,
|
||||
is_stream: bool = False,
|
||||
capability_requirements: Optional[Dict[str, bool]] = None,
|
||||
) -> Tuple[bool, Optional[str], Optional[List[str]], Optional[set[str]]]:
|
||||
capability_requirements: dict[str, bool] | None = None,
|
||||
) -> tuple[bool, str | None, list[str] | None, set[str] | None]:
|
||||
"""
|
||||
检查 Provider 是否支持指定模型(可选检查流式支持和能力需求)
|
||||
|
||||
@@ -807,12 +805,12 @@ class CacheAwareScheduler:
|
||||
self,
|
||||
db: Session,
|
||||
provider: Provider,
|
||||
global_model: "GlobalModel",
|
||||
global_model: GlobalModel,
|
||||
model_name: str,
|
||||
api_format: Optional[str] = None,
|
||||
api_format: str | None = None,
|
||||
is_stream: bool = False,
|
||||
capability_requirements: Optional[Dict[str, bool]] = None,
|
||||
) -> Tuple[bool, Optional[str], Optional[List[str]], Optional[set[str]]]:
|
||||
capability_requirements: dict[str, bool] | None = None,
|
||||
) -> tuple[bool, str | None, list[str] | None, set[str] | None]:
|
||||
"""
|
||||
检查 Provider 是否支持指定的 GlobalModel
|
||||
|
||||
@@ -841,7 +839,7 @@ class CacheAwareScheduler:
|
||||
pass
|
||||
|
||||
# 获取模型支持的能力列表
|
||||
model_supported_capabilities: List[str] = list(global_model.supported_capabilities or [])
|
||||
model_supported_capabilities: list[str] = list(global_model.supported_capabilities or [])
|
||||
|
||||
# 查询该 Provider 是否有实现这个 GlobalModel
|
||||
for model in provider.models:
|
||||
@@ -892,12 +890,12 @@ class CacheAwareScheduler:
|
||||
def _check_key_availability(
|
||||
self,
|
||||
key: ProviderAPIKey,
|
||||
api_format: Optional[str],
|
||||
api_format: str | None,
|
||||
model_name: str,
|
||||
capability_requirements: Optional[Dict[str, bool]] = None,
|
||||
model_mappings: Optional[List[str]] = None,
|
||||
candidate_models: Optional[set[str]] = None,
|
||||
) -> Tuple[bool, Optional[str], Optional[str]]:
|
||||
capability_requirements: dict[str, bool] | None = None,
|
||||
model_mappings: list[str] | None = None,
|
||||
candidate_models: set[str] | None = None,
|
||||
) -> tuple[bool, str | None, str | None]:
|
||||
"""
|
||||
检查 API Key 的可用性
|
||||
|
||||
@@ -971,7 +969,7 @@ class CacheAwareScheduler:
|
||||
# 因为 check_capability_match 会检查 Key 的 EXCLUSIVE 能力是否被浪费
|
||||
from src.core.key_capabilities import check_capability_match
|
||||
|
||||
key_caps: Dict[str, bool] = dict(key.capabilities or {})
|
||||
key_caps: dict[str, bool] = dict(key.capabilities or {})
|
||||
is_match, skip_reason = check_capability_match(key_caps, capability_requirements)
|
||||
if not is_match:
|
||||
return False, skip_reason, None
|
||||
@@ -981,16 +979,16 @@ class CacheAwareScheduler:
|
||||
async def _build_candidates(
|
||||
self,
|
||||
db: Session,
|
||||
providers: List[Provider],
|
||||
providers: list[Provider],
|
||||
client_format: APIFormat,
|
||||
model_name: str,
|
||||
affinity_key: Optional[str],
|
||||
model_mappings: Optional[List[str]] = None,
|
||||
max_candidates: Optional[int] = None,
|
||||
affinity_key: str | None,
|
||||
model_mappings: list[str] | None = None,
|
||||
max_candidates: int | None = None,
|
||||
is_stream: bool = False,
|
||||
capability_requirements: Optional[Dict[str, bool]] = None,
|
||||
capability_requirements: dict[str, bool] | None = None,
|
||||
global_conversion_enabled: bool = False,
|
||||
) -> List[ProviderCandidate]:
|
||||
) -> list[ProviderCandidate]:
|
||||
"""
|
||||
构建候选列表
|
||||
|
||||
@@ -1013,18 +1011,18 @@ class CacheAwareScheduler:
|
||||
"""
|
||||
from src.core.api_format.conversion.compatibility import is_format_compatible
|
||||
|
||||
candidates: List[ProviderCandidate] = []
|
||||
candidates: list[ProviderCandidate] = []
|
||||
client_format_str = client_format.value
|
||||
|
||||
for provider in providers:
|
||||
# 按端点格式分别判断兼容性与模型/Key 可用性:
|
||||
# - 同格式端点优先(needs_conversion=False)
|
||||
# - 跨格式端点次之(needs_conversion=True)
|
||||
model_support_cache: Dict[
|
||||
str, Tuple[bool, Optional[str], Optional[List[str]], Optional[Set[str]]]
|
||||
model_support_cache: dict[
|
||||
str, tuple[bool, str | None, list[str] | None, set[str] | None]
|
||||
] = {}
|
||||
exact_candidates: List[ProviderCandidate] = []
|
||||
convertible_candidates: List[ProviderCandidate] = []
|
||||
exact_candidates: list[ProviderCandidate] = []
|
||||
convertible_candidates: list[ProviderCandidate] = []
|
||||
|
||||
for endpoint in provider.endpoints:
|
||||
if not endpoint.is_active:
|
||||
@@ -1130,11 +1128,11 @@ class CacheAwareScheduler:
|
||||
|
||||
async def _apply_cache_affinity(
|
||||
self,
|
||||
candidates: List[ProviderCandidate],
|
||||
candidates: list[ProviderCandidate],
|
||||
affinity_key: str,
|
||||
api_format: APIFormat,
|
||||
global_model_id: str,
|
||||
) -> List[ProviderCandidate]:
|
||||
) -> list[ProviderCandidate]:
|
||||
"""
|
||||
应用缓存亲和性排序
|
||||
|
||||
@@ -1177,7 +1175,7 @@ class CacheAwareScheduler:
|
||||
return True # 需要降级
|
||||
|
||||
# 按是否匹配缓存亲和性分类候选,同时记录是否降级
|
||||
matched_candidate: Optional[ProviderCandidate] = None
|
||||
matched_candidate: ProviderCandidate | None = None
|
||||
matched = False
|
||||
|
||||
for candidate in candidates:
|
||||
@@ -1231,8 +1229,8 @@ class CacheAwareScheduler:
|
||||
matched_should_demote = should_demote(matched_candidate)
|
||||
|
||||
# 分组:非降级类 和 降级类
|
||||
keep_priority_candidates: List[ProviderCandidate] = []
|
||||
demote_candidates: List[ProviderCandidate] = []
|
||||
keep_priority_candidates: list[ProviderCandidate] = []
|
||||
demote_candidates: list[ProviderCandidate] = []
|
||||
|
||||
for c in candidates:
|
||||
if c is matched_candidate:
|
||||
@@ -1258,7 +1256,7 @@ class CacheAwareScheduler:
|
||||
logger.warning(f"检查缓存亲和性失败: {e},继续使用默认排序")
|
||||
return candidates
|
||||
|
||||
def _normalize_priority_mode(self, mode: Optional[str]) -> str:
|
||||
def _normalize_priority_mode(self, mode: str | None) -> str:
|
||||
normalized = (mode or "").strip().lower()
|
||||
if normalized not in self.ALLOWED_PRIORITY_MODES:
|
||||
if normalized:
|
||||
@@ -1266,7 +1264,7 @@ class CacheAwareScheduler:
|
||||
return self.PRIORITY_MODE_PROVIDER
|
||||
return normalized
|
||||
|
||||
def set_priority_mode(self, mode: Optional[str]) -> None:
|
||||
def set_priority_mode(self, mode: str | None) -> None:
|
||||
"""运行时更新候选排序策略"""
|
||||
normalized = self._normalize_priority_mode(mode)
|
||||
if normalized == self.priority_mode:
|
||||
@@ -1274,7 +1272,7 @@ class CacheAwareScheduler:
|
||||
self.priority_mode = normalized
|
||||
logger.debug(f"[CacheAwareScheduler] 切换优先级模式为: {self.priority_mode}")
|
||||
|
||||
def _normalize_scheduling_mode(self, mode: Optional[str]) -> str:
|
||||
def _normalize_scheduling_mode(self, mode: str | None) -> str:
|
||||
normalized = (mode or "").strip().lower()
|
||||
if normalized not in self.ALLOWED_SCHEDULING_MODES:
|
||||
if normalized:
|
||||
@@ -1284,7 +1282,7 @@ class CacheAwareScheduler:
|
||||
return self.SCHEDULING_MODE_CACHE_AFFINITY
|
||||
return normalized
|
||||
|
||||
def set_scheduling_mode(self, mode: Optional[str]) -> None:
|
||||
def set_scheduling_mode(self, mode: str | None) -> None:
|
||||
"""运行时更新调度模式"""
|
||||
normalized = self._normalize_scheduling_mode(mode)
|
||||
if normalized == self.scheduling_mode:
|
||||
@@ -1294,10 +1292,10 @@ class CacheAwareScheduler:
|
||||
|
||||
def _apply_priority_mode_sort(
|
||||
self,
|
||||
candidates: List[ProviderCandidate],
|
||||
affinity_key: Optional[str] = None,
|
||||
api_format: Optional[str] = None,
|
||||
) -> List[ProviderCandidate]:
|
||||
candidates: list[ProviderCandidate],
|
||||
affinity_key: str | None = None,
|
||||
api_format: str | None = None,
|
||||
) -> list[ProviderCandidate]:
|
||||
"""
|
||||
根据优先级模式对候选列表排序(数字越小越优先)
|
||||
|
||||
@@ -1328,8 +1326,8 @@ class CacheAwareScheduler:
|
||||
# 全局未开启:按是否需要降级分组
|
||||
# - 不需要降级:exact 候选 或 provider.keep_priority_on_conversion=True 的 convertible 候选
|
||||
# - 需要降级:convertible 且 provider.keep_priority_on_conversion=False
|
||||
keep_priority_candidates: List[ProviderCandidate] = []
|
||||
demote_candidates: List[ProviderCandidate] = []
|
||||
keep_priority_candidates: list[ProviderCandidate] = []
|
||||
demote_candidates: list[ProviderCandidate] = []
|
||||
|
||||
for c in candidates:
|
||||
if not c.needs_conversion:
|
||||
@@ -1357,10 +1355,10 @@ class CacheAwareScheduler:
|
||||
|
||||
def _sort_by_global_priority_with_hash(
|
||||
self,
|
||||
candidates: List[ProviderCandidate],
|
||||
affinity_key: Optional[str] = None,
|
||||
api_format: Optional[str] = None,
|
||||
) -> List[ProviderCandidate]:
|
||||
candidates: list[ProviderCandidate],
|
||||
affinity_key: str | None = None,
|
||||
api_format: str | None = None,
|
||||
) -> list[ProviderCandidate]:
|
||||
"""
|
||||
按 global_priority_by_format 分组排序,同优先级内通过哈希分散实现负载均衡
|
||||
|
||||
@@ -1382,7 +1380,7 @@ class CacheAwareScheduler:
|
||||
return 999999 # NULL 排在后面
|
||||
|
||||
# 按优先级分组
|
||||
priority_groups: Dict[int, List[ProviderCandidate]] = defaultdict(list)
|
||||
priority_groups: dict[int, list[ProviderCandidate]] = defaultdict(list)
|
||||
for candidate in candidates:
|
||||
priority = get_priority(candidate)
|
||||
priority_groups[priority].append(candidate)
|
||||
@@ -1417,8 +1415,8 @@ class CacheAwareScheduler:
|
||||
return result
|
||||
|
||||
def _apply_load_balance(
|
||||
self, candidates: List[ProviderCandidate], api_format: Optional[str] = None
|
||||
) -> List[ProviderCandidate]:
|
||||
self, candidates: list[ProviderCandidate], api_format: str | None = None
|
||||
) -> list[ProviderCandidate]:
|
||||
"""
|
||||
负载均衡模式:同优先级内随机轮换
|
||||
|
||||
@@ -1432,7 +1430,7 @@ class CacheAwareScheduler:
|
||||
|
||||
from collections import defaultdict
|
||||
|
||||
priority_groups: Dict[tuple, List[ProviderCandidate]] = defaultdict(list)
|
||||
priority_groups: dict[tuple, list[ProviderCandidate]] = defaultdict(list)
|
||||
|
||||
# 根据优先级模式选择分组方式
|
||||
if self.priority_mode == self.PRIORITY_MODE_GLOBAL_KEY:
|
||||
@@ -1453,7 +1451,7 @@ class CacheAwareScheduler:
|
||||
)
|
||||
priority_groups[key].append(candidate)
|
||||
|
||||
result: List[ProviderCandidate] = []
|
||||
result: list[ProviderCandidate] = []
|
||||
for priority in sorted(priority_groups.keys()):
|
||||
group = priority_groups[priority]
|
||||
if len(group) > 1:
|
||||
@@ -1468,10 +1466,10 @@ class CacheAwareScheduler:
|
||||
|
||||
def _shuffle_keys_by_internal_priority(
|
||||
self,
|
||||
keys: List[ProviderAPIKey],
|
||||
affinity_key: Optional[str] = None,
|
||||
keys: list[ProviderAPIKey],
|
||||
affinity_key: str | None = None,
|
||||
use_random: bool = False,
|
||||
) -> List[ProviderAPIKey]:
|
||||
) -> list[ProviderAPIKey]:
|
||||
"""
|
||||
对 API Key 按 internal_priority 分组,同优先级内部基于 affinity_key 进行确定性打乱
|
||||
|
||||
@@ -1495,7 +1493,7 @@ class CacheAwareScheduler:
|
||||
# 按 internal_priority 分组
|
||||
from collections import defaultdict
|
||||
|
||||
priority_groups: Dict[int, List[ProviderAPIKey]] = defaultdict(list)
|
||||
priority_groups: dict[int, list[ProviderAPIKey]] = defaultdict(list)
|
||||
|
||||
for key in keys:
|
||||
priority = key.internal_priority if key.internal_priority is not None else 999999
|
||||
@@ -1538,9 +1536,9 @@ class CacheAwareScheduler:
|
||||
affinity_key: str,
|
||||
api_format: str,
|
||||
global_model_id: str,
|
||||
endpoint_id: Optional[str] = None,
|
||||
key_id: Optional[str] = None,
|
||||
provider_id: Optional[str] = None,
|
||||
endpoint_id: str | None = None,
|
||||
key_id: str | None = None,
|
||||
provider_id: str | None = None,
|
||||
):
|
||||
"""
|
||||
失效指定亲和性标识符对特定API格式和模型的缓存亲和性
|
||||
@@ -1571,7 +1569,7 @@ class CacheAwareScheduler:
|
||||
key_id: str,
|
||||
api_format: str,
|
||||
global_model_id: str,
|
||||
ttl: Optional[int] = None,
|
||||
ttl: int | None = None,
|
||||
):
|
||||
"""
|
||||
记录缓存亲和性(供编排器调用)
|
||||
@@ -1641,13 +1639,13 @@ class CacheAwareScheduler:
|
||||
|
||||
|
||||
# 全局单例
|
||||
_scheduler: Optional[CacheAwareScheduler] = None
|
||||
_scheduler: CacheAwareScheduler | None = None
|
||||
|
||||
|
||||
async def get_cache_aware_scheduler(
|
||||
redis_client=None,
|
||||
priority_mode: Optional[str] = None,
|
||||
scheduling_mode: Optional[str] = None,
|
||||
priority_mode: str | None = None,
|
||||
scheduling_mode: str | None = None,
|
||||
) -> CacheAwareScheduler:
|
||||
"""
|
||||
获取全局CacheAwareScheduler实例
|
||||
|
||||
22
src/services/cache/backend.py
vendored
22
src/services/cache/backend.py
vendored
@@ -15,7 +15,7 @@ import json
|
||||
import time
|
||||
from abc import ABC, abstractmethod
|
||||
from collections import OrderedDict
|
||||
from typing import Any, Dict, Optional
|
||||
from typing import Any
|
||||
|
||||
import redis.asyncio as aioredis
|
||||
from src.core.logger import logger
|
||||
@@ -28,7 +28,7 @@ class BaseCacheBackend(ABC):
|
||||
"""缓存后端抽象基类"""
|
||||
|
||||
@abstractmethod
|
||||
async def get(self, key: str) -> Optional[Any]:
|
||||
async def get(self, key: str) -> Any | None:
|
||||
"""获取缓存值"""
|
||||
pass
|
||||
|
||||
@@ -43,7 +43,7 @@ class BaseCacheBackend(ABC):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def clear(self, pattern: Optional[str] = None) -> None:
|
||||
async def clear(self, pattern: str | None = None) -> None:
|
||||
"""清空缓存(支持模式匹配)"""
|
||||
pass
|
||||
|
||||
@@ -65,12 +65,12 @@ class LocalCache(BaseCacheBackend):
|
||||
default_ttl: 默认过期时间(秒)
|
||||
"""
|
||||
self._cache: OrderedDict = OrderedDict()
|
||||
self._expiry: Dict[str, float] = {}
|
||||
self._expiry: dict[str, float] = {}
|
||||
self._max_size = max_size
|
||||
self._default_ttl = default_ttl
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
async def get(self, key: str) -> Optional[Any]:
|
||||
async def get(self, key: str) -> Any | None:
|
||||
"""获取缓存值(线程安全)"""
|
||||
async with self._lock:
|
||||
if key not in self._cache:
|
||||
@@ -115,7 +115,7 @@ class LocalCache(BaseCacheBackend):
|
||||
if key in self._expiry:
|
||||
del self._expiry[key]
|
||||
|
||||
async def clear(self, pattern: Optional[str] = None) -> None:
|
||||
async def clear(self, pattern: str | None = None) -> None:
|
||||
"""清空缓存(线程安全)"""
|
||||
async with self._lock:
|
||||
if pattern is None:
|
||||
@@ -145,7 +145,7 @@ class LocalCache(BaseCacheBackend):
|
||||
|
||||
return True
|
||||
|
||||
def get_stats(self) -> Dict[str, Any]:
|
||||
def get_stats(self) -> dict[str, Any]:
|
||||
"""获取缓存统计信息"""
|
||||
return {
|
||||
"backend": "local",
|
||||
@@ -177,7 +177,7 @@ class RedisCache(BaseCacheBackend):
|
||||
"""构造完整的 Redis 键"""
|
||||
return f"{self._key_prefix}:{key}"
|
||||
|
||||
async def get(self, key: str) -> Optional[Any]:
|
||||
async def get(self, key: str) -> Any | None:
|
||||
"""获取缓存值"""
|
||||
try:
|
||||
redis_key = self._make_key(key)
|
||||
@@ -223,7 +223,7 @@ class RedisCache(BaseCacheBackend):
|
||||
except Exception as e:
|
||||
logger.error(f"[RedisCache] 删除缓存失败: {key}, 错误: {e}")
|
||||
|
||||
async def clear(self, pattern: Optional[str] = None) -> None:
|
||||
async def clear(self, pattern: str | None = None) -> None:
|
||||
"""清空缓存"""
|
||||
try:
|
||||
if pattern is None:
|
||||
@@ -264,7 +264,7 @@ class RedisCache(BaseCacheBackend):
|
||||
except Exception as e:
|
||||
logger.error(f"[RedisCache] 发布缓存失效失败: {channel}, {key}, 错误: {e}")
|
||||
|
||||
def get_stats(self) -> Dict[str, Any]:
|
||||
def get_stats(self) -> dict[str, Any]:
|
||||
"""获取缓存统计信息"""
|
||||
return {
|
||||
"backend": "redis",
|
||||
@@ -274,7 +274,7 @@ class RedisCache(BaseCacheBackend):
|
||||
|
||||
|
||||
# 缓存后端工厂
|
||||
_cache_backends: Dict[str, BaseCacheBackend] = {}
|
||||
_cache_backends: dict[str, BaseCacheBackend] = {}
|
||||
|
||||
|
||||
async def get_cache_backend(
|
||||
|
||||
5
src/services/cache/invalidation.py
vendored
5
src/services/cache/invalidation.py
vendored
@@ -4,7 +4,6 @@
|
||||
统一管理各种缓存的失效逻辑
|
||||
"""
|
||||
|
||||
from typing import Optional
|
||||
|
||||
from src.core.logger import logger
|
||||
|
||||
@@ -21,7 +20,7 @@ class CacheInvalidationService:
|
||||
self._model_mappers.append(model_mapper)
|
||||
|
||||
async def on_global_model_changed(
|
||||
self, model_name: str, global_model_id: Optional[str] = None
|
||||
self, model_name: str, global_model_id: str | None = None
|
||||
) -> None:
|
||||
"""
|
||||
GlobalModel 变更时的缓存失效
|
||||
@@ -96,7 +95,7 @@ class CacheInvalidationService:
|
||||
|
||||
|
||||
# 全局单例
|
||||
_cache_invalidation_service: Optional[CacheInvalidationService] = None
|
||||
_cache_invalidation_service: CacheInvalidationService | None = None
|
||||
|
||||
|
||||
def get_cache_invalidation_service() -> CacheInvalidationService:
|
||||
|
||||
27
src/services/cache/model_cache.py
vendored
27
src/services/cache/model_cache.py
vendored
@@ -19,7 +19,6 @@ Model 映射缓存服务 - 减少模型查询
|
||||
"""
|
||||
|
||||
import time
|
||||
from typing import List, Optional
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
@@ -45,7 +44,7 @@ class ModelCacheService:
|
||||
CACHE_TTL = CacheTTL.MODEL
|
||||
|
||||
@staticmethod
|
||||
async def get_model_by_id(db: Session, model_id: str) -> Optional[Model]:
|
||||
async def get_model_by_id(db: Session, model_id: str) -> Model | None:
|
||||
"""
|
||||
获取 Model(带缓存)
|
||||
|
||||
@@ -76,7 +75,7 @@ class ModelCacheService:
|
||||
return model
|
||||
|
||||
@staticmethod
|
||||
async def get_global_model_by_id(db: Session, global_model_id: str) -> Optional[GlobalModel]:
|
||||
async def get_global_model_by_id(db: Session, global_model_id: str) -> GlobalModel | None:
|
||||
"""
|
||||
获取 GlobalModel(带缓存)
|
||||
|
||||
@@ -111,7 +110,7 @@ class ModelCacheService:
|
||||
@staticmethod
|
||||
async def get_model_by_provider_and_global_model(
|
||||
db: Session, provider_id: str, global_model_id: str
|
||||
) -> Optional[Model]:
|
||||
) -> Model | None:
|
||||
"""
|
||||
通过 Provider ID 和 GlobalModel ID 获取 Model(带缓存)
|
||||
|
||||
@@ -160,7 +159,7 @@ class ModelCacheService:
|
||||
return model
|
||||
|
||||
@staticmethod
|
||||
async def get_global_model_by_name(db: Session, name: str) -> Optional[GlobalModel]:
|
||||
async def get_global_model_by_name(db: Session, name: str) -> GlobalModel | None:
|
||||
"""
|
||||
通过名称获取 GlobalModel(带缓存)
|
||||
|
||||
@@ -195,10 +194,10 @@ class ModelCacheService:
|
||||
@staticmethod
|
||||
async def invalidate_model_cache(
|
||||
model_id: str,
|
||||
provider_id: Optional[str] = None,
|
||||
global_model_id: Optional[str] = None,
|
||||
provider_model_name: Optional[str] = None,
|
||||
provider_model_mappings: Optional[list] = None,
|
||||
provider_id: str | None = None,
|
||||
global_model_id: str | None = None,
|
||||
provider_model_name: str | None = None,
|
||||
provider_model_mappings: list | None = None,
|
||||
) -> None:
|
||||
"""清除 Model 缓存
|
||||
|
||||
@@ -240,7 +239,7 @@ class ModelCacheService:
|
||||
logger.debug(f"Model resolve 缓存已清除: {resolve_keys_to_clear}")
|
||||
|
||||
@staticmethod
|
||||
async def invalidate_global_model_cache(global_model_id: str, name: Optional[str] = None) -> None:
|
||||
async def invalidate_global_model_cache(global_model_id: str, name: str | None = None) -> None:
|
||||
"""清除 GlobalModel 缓存"""
|
||||
await CacheService.delete(f"global_model:id:{global_model_id}")
|
||||
if name:
|
||||
@@ -270,7 +269,7 @@ class ModelCacheService:
|
||||
@staticmethod
|
||||
async def resolve_global_model_by_name_or_mapping(
|
||||
db: Session, model_name: str
|
||||
) -> Optional[GlobalModel]:
|
||||
) -> GlobalModel | None:
|
||||
"""
|
||||
通过名称解析 GlobalModel(带缓存)
|
||||
|
||||
@@ -355,7 +354,7 @@ class ModelCacheService:
|
||||
)
|
||||
|
||||
# 收集匹配的 GlobalModel(只通过 provider_model_name 匹配)
|
||||
matched_global_models: List[GlobalModel] = []
|
||||
matched_global_models: list[GlobalModel] = []
|
||||
seen_global_model_ids: set[str] = set()
|
||||
for model, gm in models_with_global:
|
||||
if gm.id not in seen_global_model_ids:
|
||||
@@ -405,7 +404,7 @@ class ModelCacheService:
|
||||
.all()
|
||||
)
|
||||
|
||||
mapping_matched_global_models: List[GlobalModel] = []
|
||||
mapping_matched_global_models: list[GlobalModel] = []
|
||||
mapping_seen_ids: set[str] = set()
|
||||
for model, gm in models_with_mappings:
|
||||
raw_mappings = model.provider_model_mappings
|
||||
@@ -469,7 +468,7 @@ class ModelCacheService:
|
||||
.all()
|
||||
)
|
||||
|
||||
mapping_matches: List[GlobalModel] = []
|
||||
mapping_matches: list[GlobalModel] = []
|
||||
for gm in mapping_rows:
|
||||
config = gm.config or {}
|
||||
mappings = config.get("model_mappings")
|
||||
|
||||
19
src/services/cache/provider_cache.py
vendored
19
src/services/cache/provider_cache.py
vendored
@@ -5,7 +5,6 @@ Provider 缓存服务 - 减少 Provider 和 ProviderAPIKey 查询
|
||||
这些数据在 UsageService.record_usage() 中被频繁查询但变化不频繁。
|
||||
"""
|
||||
|
||||
from typing import Optional, Tuple
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
@@ -27,8 +26,8 @@ class ProviderCacheService:
|
||||
|
||||
@staticmethod
|
||||
def compute_rate_multiplier(
|
||||
rate_multipliers: Optional[dict],
|
||||
api_format: Optional[str] = None,
|
||||
rate_multipliers: dict | None,
|
||||
api_format: str | None = None,
|
||||
) -> float:
|
||||
"""
|
||||
计算 rate_multiplier 的纯函数(无数据库/缓存依赖)
|
||||
@@ -50,8 +49,8 @@ class ProviderCacheService:
|
||||
|
||||
@staticmethod
|
||||
async def get_provider_api_key_rate_multiplier(
|
||||
db: Session, provider_api_key_id: str, api_format: Optional[str] = None
|
||||
) -> Optional[float]:
|
||||
db: Session, provider_api_key_id: str, api_format: str | None = None
|
||||
) -> float | None:
|
||||
"""
|
||||
获取 ProviderAPIKey 的 rate_multiplier(带缓存)
|
||||
|
||||
@@ -106,7 +105,7 @@ class ProviderCacheService:
|
||||
@staticmethod
|
||||
async def get_provider_billing_type(
|
||||
db: Session, provider_id: str
|
||||
) -> Optional[ProviderBillingType]:
|
||||
) -> ProviderBillingType | None:
|
||||
"""
|
||||
获取 Provider 的 billing_type(带缓存)
|
||||
|
||||
@@ -154,10 +153,10 @@ class ProviderCacheService:
|
||||
@staticmethod
|
||||
async def get_rate_multiplier_and_free_tier(
|
||||
db: Session,
|
||||
provider_api_key_id: Optional[str],
|
||||
provider_id: Optional[str],
|
||||
api_format: Optional[str] = None,
|
||||
) -> Tuple[float, bool]:
|
||||
provider_api_key_id: str | None,
|
||||
provider_id: str | None,
|
||||
api_format: str | None = None,
|
||||
) -> tuple[float, bool]:
|
||||
"""
|
||||
获取费率倍数和是否免费套餐(带缓存)
|
||||
|
||||
|
||||
13
src/services/cache/sync.py
vendored
13
src/services/cache/sync.py
vendored
@@ -11,7 +11,8 @@
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from typing import Callable, Dict, Optional
|
||||
|
||||
from collections.abc import Callable
|
||||
|
||||
import redis.asyncio as aioredis
|
||||
from src.core.logger import logger
|
||||
@@ -40,9 +41,9 @@ class CacheSyncService:
|
||||
redis_client: Redis 客户端实例
|
||||
"""
|
||||
self._redis = redis_client
|
||||
self._pubsub: Optional[aioredis.client.PubSub] = None
|
||||
self._listener_task: Optional[asyncio.Task] = None
|
||||
self._handlers: Dict[str, Callable] = {}
|
||||
self._pubsub: aioredis.client.PubSub | None = None
|
||||
self._listener_task: asyncio.Task | None = None
|
||||
self._handlers: dict[str, Callable] = {}
|
||||
self._running = False
|
||||
|
||||
async def start(self):
|
||||
@@ -160,10 +161,10 @@ class CacheSyncService:
|
||||
|
||||
|
||||
# 全局单例
|
||||
_cache_sync_service: Optional[CacheSyncService] = None
|
||||
_cache_sync_service: CacheSyncService | None = None
|
||||
|
||||
|
||||
async def get_cache_sync_service(redis_client: aioredis.Redis = None) -> Optional[CacheSyncService]:
|
||||
async def get_cache_sync_service(redis_client: aioredis.Redis = None) -> CacheSyncService | None:
|
||||
"""
|
||||
获取缓存同步服务实例
|
||||
|
||||
|
||||
7
src/services/cache/user_cache.py
vendored
7
src/services/cache/user_cache.py
vendored
@@ -19,7 +19,6 @@
|
||||
await UserCacheService.invalidate_user_cache(user_id, email)
|
||||
"""
|
||||
|
||||
from typing import Optional
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
@@ -40,7 +39,7 @@ class UserCacheService:
|
||||
CACHE_TTL = CacheTTL.USER
|
||||
|
||||
@staticmethod
|
||||
async def get_user_by_id(db: Session, user_id: str) -> Optional[User]:
|
||||
async def get_user_by_id(db: Session, user_id: str) -> User | None:
|
||||
"""
|
||||
获取用户(带缓存)
|
||||
|
||||
@@ -72,7 +71,7 @@ class UserCacheService:
|
||||
return user
|
||||
|
||||
@staticmethod
|
||||
async def get_user_by_email(db: Session, email: str) -> Optional[User]:
|
||||
async def get_user_by_email(db: Session, email: str) -> User | None:
|
||||
"""
|
||||
通过邮箱获取用户(带缓存)
|
||||
|
||||
@@ -103,7 +102,7 @@ class UserCacheService:
|
||||
return user
|
||||
|
||||
@staticmethod
|
||||
async def invalidate_user_cache(user_id: str, email: Optional[str] = None):
|
||||
async def invalidate_user_cache(user_id: str, email: str | None = None):
|
||||
"""
|
||||
清除用户缓存
|
||||
|
||||
|
||||
@@ -9,7 +9,9 @@
|
||||
5. 显式传入 (用于重试升级)
|
||||
"""
|
||||
|
||||
from typing import Any, Callable, Dict, Optional
|
||||
from typing import Any
|
||||
|
||||
from collections.abc import Callable
|
||||
|
||||
from src.core.key_capabilities import (
|
||||
CAPABILITY_DEFINITIONS,
|
||||
@@ -20,7 +22,7 @@ from src.core.api_format import get_header_value
|
||||
from src.core.logger import logger
|
||||
|
||||
# Adapter 检测器类型:接受 headers 和可选的 request_body,返回能力需求字典
|
||||
AdapterDetectorType = Callable[[Dict[str, str], Optional[Dict[str, Any]]], Dict[str, bool]]
|
||||
type AdapterDetectorType = Callable[[dict[str, str], dict[str, Any] | None], dict[str, bool]]
|
||||
|
||||
|
||||
class CapabilityResolver:
|
||||
@@ -28,14 +30,14 @@ class CapabilityResolver:
|
||||
|
||||
@staticmethod
|
||||
def resolve_requirements(
|
||||
user: Optional[Any] = None,
|
||||
user_api_key: Optional[Any] = None,
|
||||
model_name: Optional[str] = None,
|
||||
request_headers: Optional[Dict[str, str]] = None,
|
||||
request_body: Optional[Dict[str, Any]] = None,
|
||||
explicit_requirements: Optional[Dict[str, bool]] = None,
|
||||
adapter_detector: Optional[AdapterDetectorType] = None,
|
||||
) -> Dict[str, bool]:
|
||||
user: Any | None = None,
|
||||
user_api_key: Any | None = None,
|
||||
model_name: str | None = None,
|
||||
request_headers: dict[str, str] | None = None,
|
||||
request_body: dict[str, Any] | None = None,
|
||||
explicit_requirements: dict[str, bool] | None = None,
|
||||
adapter_detector: AdapterDetectorType | None = None,
|
||||
) -> dict[str, bool]:
|
||||
"""
|
||||
解析请求的能力需求
|
||||
|
||||
@@ -58,7 +60,7 @@ class CapabilityResolver:
|
||||
Returns:
|
||||
能力需求字典,如 {"cache_1h": True, "context_1m": False}
|
||||
"""
|
||||
requirements: Dict[str, bool] = {}
|
||||
requirements: dict[str, bool] = {}
|
||||
|
||||
# 1. 从用户模型级配置获取(仅用户可配置型能力)
|
||||
if user and model_name:
|
||||
@@ -125,9 +127,9 @@ class CapabilityResolver:
|
||||
|
||||
@staticmethod
|
||||
def get_default_requirements_for_model(
|
||||
user: Optional[Any] = None,
|
||||
model_name: Optional[str] = None,
|
||||
) -> Dict[str, bool]:
|
||||
user: Any | None = None,
|
||||
model_name: str | None = None,
|
||||
) -> dict[str, bool]:
|
||||
"""
|
||||
获取用户对特定模型的默认能力需求
|
||||
|
||||
@@ -140,7 +142,7 @@ class CapabilityResolver:
|
||||
Returns:
|
||||
能力需求字典
|
||||
"""
|
||||
requirements: Dict[str, bool] = {}
|
||||
requirements: dict[str, bool] = {}
|
||||
|
||||
if not user or not model_name:
|
||||
return requirements
|
||||
@@ -156,9 +158,9 @@ class CapabilityResolver:
|
||||
|
||||
@staticmethod
|
||||
def merge_requirements(
|
||||
base: Optional[Dict[str, bool]],
|
||||
override: Optional[Dict[str, bool]],
|
||||
) -> Dict[str, bool]:
|
||||
base: dict[str, bool] | None,
|
||||
override: dict[str, bool] | None,
|
||||
) -> dict[str, bool]:
|
||||
"""
|
||||
合并两个能力需求字典
|
||||
|
||||
|
||||
@@ -6,7 +6,7 @@
|
||||
import smtplib
|
||||
from email.mime.multipart import MIMEMultipart
|
||||
from email.mime.text import MIMEText
|
||||
from typing import Any, Optional, Tuple, Union
|
||||
from typing import Any
|
||||
|
||||
aiosmtplib: Any
|
||||
try:
|
||||
@@ -69,7 +69,7 @@ class EmailSenderService:
|
||||
return config
|
||||
|
||||
@staticmethod
|
||||
def _validate_smtp_config(config: dict) -> Tuple[bool, Optional[str]]:
|
||||
def _validate_smtp_config(config: dict) -> tuple[bool, str | None]:
|
||||
"""
|
||||
验证 SMTP 配置
|
||||
|
||||
@@ -105,7 +105,7 @@ class EmailSenderService:
|
||||
@staticmethod
|
||||
async def send_verification_code(
|
||||
db: Session, to_email: str, code: str, expire_minutes: int = 30
|
||||
) -> Tuple[bool, Optional[str]]:
|
||||
) -> tuple[bool, str | None]:
|
||||
"""
|
||||
发送验证码邮件
|
||||
|
||||
@@ -151,9 +151,9 @@ class EmailSenderService:
|
||||
config: dict,
|
||||
to_email: str,
|
||||
subject: str,
|
||||
html_body: Optional[str] = None,
|
||||
text_body: Optional[str] = None,
|
||||
) -> Tuple[bool, Optional[str]]:
|
||||
html_body: str | None = None,
|
||||
text_body: str | None = None,
|
||||
) -> tuple[bool, str | None]:
|
||||
"""
|
||||
发送邮件(内部方法)
|
||||
|
||||
@@ -181,9 +181,9 @@ class EmailSenderService:
|
||||
config: dict,
|
||||
to_email: str,
|
||||
subject: str,
|
||||
html_body: Optional[str] = None,
|
||||
text_body: Optional[str] = None,
|
||||
) -> Tuple[bool, Optional[str]]:
|
||||
html_body: str | None = None,
|
||||
text_body: str | None = None,
|
||||
) -> tuple[bool, str | None]:
|
||||
"""
|
||||
异步发送邮件(使用 aiosmtplib)
|
||||
|
||||
@@ -250,9 +250,9 @@ class EmailSenderService:
|
||||
config: dict,
|
||||
to_email: str,
|
||||
subject: str,
|
||||
html_body: Optional[str] = None,
|
||||
text_body: Optional[str] = None,
|
||||
) -> Tuple[bool, Optional[str]]:
|
||||
html_body: str | None = None,
|
||||
text_body: str | None = None,
|
||||
) -> tuple[bool, str | None]:
|
||||
"""
|
||||
同步邮件发送的异步包装器
|
||||
|
||||
@@ -280,9 +280,9 @@ class EmailSenderService:
|
||||
config: dict,
|
||||
to_email: str,
|
||||
subject: str,
|
||||
html_body: Optional[str] = None,
|
||||
text_body: Optional[str] = None,
|
||||
) -> Tuple[bool, Optional[str]]:
|
||||
html_body: str | None = None,
|
||||
text_body: str | None = None,
|
||||
) -> tuple[bool, str | None]:
|
||||
"""
|
||||
同步发送邮件(使用标准库 smtplib)
|
||||
|
||||
@@ -312,7 +312,7 @@ class EmailSenderService:
|
||||
message.attach(MIMEText(html_body, "html", "utf-8"))
|
||||
|
||||
# 连接 SMTP 服务器
|
||||
server: Optional[smtplib.SMTP] = None
|
||||
server: smtplib.SMTP | None = None
|
||||
ssl_context = get_ssl_context()
|
||||
try:
|
||||
if config["smtp_use_ssl"]:
|
||||
@@ -357,8 +357,8 @@ class EmailSenderService:
|
||||
|
||||
@staticmethod
|
||||
async def test_smtp_connection(
|
||||
db: Session, override_config: Optional[dict] = None
|
||||
) -> Tuple[bool, Optional[str]]:
|
||||
db: Session, override_config: dict | None = None
|
||||
) -> tuple[bool, str | None]:
|
||||
"""
|
||||
测试 SMTP 连接
|
||||
|
||||
@@ -405,7 +405,7 @@ class EmailSenderService:
|
||||
await smtp.quit()
|
||||
else:
|
||||
# 使用同步方式测试
|
||||
server: Union[smtplib.SMTP, smtplib.SMTP_SSL]
|
||||
server: smtplib.SMTP | smtplib.SMTP_SSL
|
||||
if config["smtp_use_ssl"]:
|
||||
server = smtplib.SMTP_SSL(
|
||||
config["smtp_host"],
|
||||
|
||||
@@ -6,7 +6,7 @@
|
||||
import html
|
||||
import re
|
||||
from html.parser import HTMLParser
|
||||
from typing import Any, Dict, Optional
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
@@ -219,7 +219,7 @@ class EmailTemplate:
|
||||
</html>"""
|
||||
|
||||
@staticmethod
|
||||
def get_default_template(template_type: str) -> Dict[str, str]:
|
||||
def get_default_template(template_type: str) -> dict[str, str]:
|
||||
"""
|
||||
获取默认模板
|
||||
|
||||
@@ -243,7 +243,7 @@ class EmailTemplate:
|
||||
return {"subject": "通知", "html": ""}
|
||||
|
||||
@staticmethod
|
||||
def get_template(db: Session, template_type: str) -> Dict[str, str]:
|
||||
def get_template(db: Session, template_type: str) -> dict[str, str]:
|
||||
"""
|
||||
从数据库获取模板,如果不存在则返回默认模板
|
||||
|
||||
@@ -269,7 +269,7 @@ class EmailTemplate:
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def render_template(template_html: str, variables: Dict[str, Any]) -> str:
|
||||
def render_template(template_html: str, variables: dict[str, Any]) -> str:
|
||||
"""
|
||||
渲染模板,替换 {{variable}} 格式的变量
|
||||
|
||||
@@ -310,7 +310,7 @@ class EmailTemplate:
|
||||
|
||||
@staticmethod
|
||||
def get_verification_code_html(
|
||||
code: str, expire_minutes: int = 5, db: Optional[Session] = None, **kwargs
|
||||
code: str, expire_minutes: int = 5, db: Session | None = None, **kwargs
|
||||
) -> str:
|
||||
"""
|
||||
获取验证码邮件 HTML
|
||||
@@ -345,7 +345,7 @@ class EmailTemplate:
|
||||
|
||||
@staticmethod
|
||||
def get_verification_code_text(
|
||||
code: str, expire_minutes: int = 5, db: Optional[Session] = None, **kwargs
|
||||
code: str, expire_minutes: int = 5, db: Session | None = None, **kwargs
|
||||
) -> str:
|
||||
"""
|
||||
获取验证码邮件纯文本(从 HTML 自动生成)
|
||||
@@ -364,7 +364,7 @@ class EmailTemplate:
|
||||
|
||||
@staticmethod
|
||||
def get_password_reset_html(
|
||||
reset_link: str, expire_minutes: int = 30, db: Optional[Session] = None, **kwargs
|
||||
reset_link: str, expire_minutes: int = 30, db: Session | None = None, **kwargs
|
||||
) -> str:
|
||||
"""
|
||||
获取密码重置邮件 HTML
|
||||
@@ -399,7 +399,7 @@ class EmailTemplate:
|
||||
|
||||
@staticmethod
|
||||
def get_password_reset_text(
|
||||
reset_link: str, expire_minutes: int = 30, db: Optional[Session] = None, **kwargs
|
||||
reset_link: str, expire_minutes: int = 30, db: Session | None = None, **kwargs
|
||||
) -> str:
|
||||
"""
|
||||
获取密码重置邮件纯文本(从 HTML 自动生成)
|
||||
@@ -418,7 +418,7 @@ class EmailTemplate:
|
||||
|
||||
@staticmethod
|
||||
def get_subject(
|
||||
template_type: str = "verification", db: Optional[Session] = None
|
||||
template_type: str = "verification", db: Session | None = None
|
||||
) -> str:
|
||||
"""
|
||||
获取邮件主题
|
||||
|
||||
@@ -6,7 +6,6 @@
|
||||
import json
|
||||
import secrets
|
||||
from datetime import datetime, timezone
|
||||
from typing import Optional, Tuple
|
||||
|
||||
from src.clients.redis_client import get_redis_client
|
||||
from src.config.settings import Config
|
||||
@@ -42,8 +41,8 @@ class EmailVerificationService:
|
||||
@staticmethod
|
||||
async def send_verification_code(
|
||||
email: str,
|
||||
expire_minutes: Optional[int] = None,
|
||||
) -> Tuple[bool, str, Optional[str]]:
|
||||
expire_minutes: int | None = None,
|
||||
) -> tuple[bool, str, str | None]:
|
||||
"""
|
||||
发送验证码(生成并存储到 Redis)
|
||||
|
||||
@@ -98,7 +97,7 @@ class EmailVerificationService:
|
||||
return False, "系统错误", str(e)
|
||||
|
||||
@staticmethod
|
||||
async def verify_code(email: str, code: str) -> Tuple[bool, str]:
|
||||
async def verify_code(email: str, code: str) -> tuple[bool, str]:
|
||||
"""
|
||||
验证验证码
|
||||
|
||||
|
||||
@@ -11,7 +11,7 @@
|
||||
import json
|
||||
from collections import defaultdict
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Any, Dict, List, Optional
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import case, func
|
||||
from sqlalchemy.orm import Session
|
||||
@@ -43,7 +43,7 @@ class EndpointHealthService:
|
||||
lookback_hours: int = 6,
|
||||
include_admin_fields: bool = False,
|
||||
use_cache: bool = True,
|
||||
) -> List[Dict[str, Any]]:
|
||||
) -> list[dict[str, Any]]:
|
||||
"""
|
||||
获取按 API 格式聚合的端点健康状态
|
||||
|
||||
@@ -71,7 +71,7 @@ class EndpointHealthService:
|
||||
)
|
||||
|
||||
# 收集所有 provider_ids
|
||||
all_provider_ids = list(set(ep.provider_id for ep in endpoints))
|
||||
all_provider_ids = list({ep.provider_id for ep in endpoints})
|
||||
|
||||
# 批量查询所有密钥(通过 provider_id 关联)
|
||||
all_keys = (
|
||||
@@ -81,7 +81,7 @@ class EndpointHealthService:
|
||||
) if all_provider_ids else []
|
||||
|
||||
# 按 api_format 分组密钥(通过 api_formats 字段)
|
||||
keys_by_format: Dict[str, List[ProviderAPIKey]] = defaultdict(list)
|
||||
keys_by_format: dict[str, list[ProviderAPIKey]] = defaultdict(list)
|
||||
for key in all_keys:
|
||||
for fmt in (key.api_formats or []):
|
||||
keys_by_format[fmt].append(key)
|
||||
@@ -140,7 +140,7 @@ class EndpointHealthService:
|
||||
|
||||
# 批量生成所有格式的时间线数据
|
||||
all_key_ids = []
|
||||
format_key_mapping: Dict[str, List[str]] = {}
|
||||
format_key_mapping: dict[str, list[str]] = {}
|
||||
for api_format, stats in format_stats.items():
|
||||
key_ids = stats["key_ids"]
|
||||
format_key_mapping[api_format] = key_ids
|
||||
@@ -215,11 +215,11 @@ class EndpointHealthService:
|
||||
@staticmethod
|
||||
def _generate_timeline_batch(
|
||||
db: Session,
|
||||
format_key_mapping: Dict[str, List[str]],
|
||||
format_key_mapping: dict[str, list[str]],
|
||||
now: datetime,
|
||||
lookback_hours: int,
|
||||
segments: int = 100,
|
||||
) -> Dict[str, Dict[str, Any]]:
|
||||
) -> dict[str, dict[str, Any]]:
|
||||
"""
|
||||
批量生成多个 API 格式的时间线数据(基于 RequestCandidate 表)
|
||||
|
||||
@@ -303,13 +303,13 @@ class EndpointHealthService:
|
||||
)
|
||||
|
||||
# 构建 key_id -> api_format 的反向映射
|
||||
key_to_format: Dict[str, str] = {}
|
||||
key_to_format: dict[str, str] = {}
|
||||
for api_format, key_ids in format_key_mapping.items():
|
||||
for key_id in key_ids:
|
||||
key_to_format[key_id] = api_format
|
||||
|
||||
# 按 api_format 和 segment 聚合数据
|
||||
format_segment_data: Dict[str, Dict[int, Dict]] = defaultdict(lambda: defaultdict(lambda: {
|
||||
format_segment_data: dict[str, dict[int, dict]] = defaultdict(lambda: defaultdict(lambda: {
|
||||
"total": 0,
|
||||
"success": 0,
|
||||
"failed": 0,
|
||||
@@ -336,7 +336,7 @@ class EndpointHealthService:
|
||||
seg_data["max_time"] = row.max_time
|
||||
|
||||
# 生成各格式的时间线
|
||||
result: Dict[str, Dict[str, Any]] = {}
|
||||
result: dict[str, dict[str, Any]] = {}
|
||||
|
||||
for api_format in format_key_mapping.keys():
|
||||
timeline = []
|
||||
@@ -385,11 +385,11 @@ class EndpointHealthService:
|
||||
@staticmethod
|
||||
def _generate_timeline_from_usage(
|
||||
db: Session,
|
||||
endpoint_ids: List[str],
|
||||
endpoint_ids: list[str],
|
||||
now: datetime,
|
||||
lookback_hours: int,
|
||||
segments: int = 100,
|
||||
) -> Dict[str, Any]:
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
从真实使用记录生成时间线数据(使用批量查询优化)
|
||||
|
||||
@@ -472,7 +472,7 @@ class EndpointHealthService:
|
||||
return format_names.get(api_format, api_format)
|
||||
|
||||
@staticmethod
|
||||
def _get_from_cache(key: str) -> Optional[List[Dict[str, Any]]]:
|
||||
def _get_from_cache(key: str) -> list[dict[str, Any]] | None:
|
||||
"""从 Redis 缓存获取数据"""
|
||||
redis_client = _get_redis_client()
|
||||
if not redis_client:
|
||||
@@ -487,7 +487,7 @@ class EndpointHealthService:
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _set_to_cache(key: str, data: List[Dict[str, Any]]) -> None:
|
||||
def _set_to_cache(key: str, data: list[dict[str, Any]]) -> None:
|
||||
"""写入 Redis 缓存"""
|
||||
redis_client = _get_redis_client()
|
||||
if not redis_client:
|
||||
|
||||
@@ -14,7 +14,7 @@
|
||||
|
||||
import os
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import case, func
|
||||
from sqlalchemy.orm import Session
|
||||
@@ -35,7 +35,7 @@ class CircuitState:
|
||||
|
||||
|
||||
# 默认健康度数据结构
|
||||
def _default_health_data() -> Dict[str, Any]:
|
||||
def _default_health_data() -> dict[str, Any]:
|
||||
return {
|
||||
"health_score": 1.0,
|
||||
"consecutive_failures": 0,
|
||||
@@ -45,7 +45,7 @@ def _default_health_data() -> Dict[str, Any]:
|
||||
|
||||
|
||||
# 默认熔断器数据结构
|
||||
def _default_circuit_data() -> Dict[str, Any]:
|
||||
def _default_circuit_data() -> dict[str, Any]:
|
||||
return {
|
||||
"open": False,
|
||||
"open_at": None,
|
||||
@@ -119,13 +119,13 @@ class HealthMonitor:
|
||||
CIRCUIT_HISTORY_LIMIT = int(os.getenv("HEALTH_CIRCUIT_HISTORY_LIMIT", "200"))
|
||||
|
||||
# 进程级别状态缓存
|
||||
_circuit_history: List[Dict[str, Any]] = []
|
||||
_circuit_history: list[dict[str, Any]] = []
|
||||
_open_circuit_keys: int = 0
|
||||
|
||||
# ==================== 数据访问辅助方法 ====================
|
||||
|
||||
@classmethod
|
||||
def _get_health_data(cls, key: ProviderAPIKey, api_format: str) -> Dict[str, Any]:
|
||||
def _get_health_data(cls, key: ProviderAPIKey, api_format: str) -> dict[str, Any]:
|
||||
"""获取指定格式的健康度数据,不存在则返回默认值"""
|
||||
health_by_format = key.health_by_format or {}
|
||||
if api_format not in health_by_format:
|
||||
@@ -133,14 +133,14 @@ class HealthMonitor:
|
||||
return health_by_format[api_format]
|
||||
|
||||
@classmethod
|
||||
def _set_health_data(cls, key: ProviderAPIKey, api_format: str, data: Dict[str, Any]) -> None:
|
||||
def _set_health_data(cls, key: ProviderAPIKey, api_format: str, data: dict[str, Any]) -> None:
|
||||
"""设置指定格式的健康度数据"""
|
||||
health_by_format = dict(key.health_by_format or {})
|
||||
health_by_format[api_format] = data
|
||||
key.health_by_format = health_by_format # type: ignore[assignment]
|
||||
|
||||
@classmethod
|
||||
def _get_circuit_data(cls, key: ProviderAPIKey, api_format: str) -> Dict[str, Any]:
|
||||
def _get_circuit_data(cls, key: ProviderAPIKey, api_format: str) -> dict[str, Any]:
|
||||
"""获取指定格式的熔断器数据,不存在则返回默认值"""
|
||||
circuit_by_format = key.circuit_breaker_by_format or {}
|
||||
if api_format not in circuit_by_format:
|
||||
@@ -148,7 +148,7 @@ class HealthMonitor:
|
||||
return circuit_by_format[api_format]
|
||||
|
||||
@classmethod
|
||||
def _set_circuit_data(cls, key: ProviderAPIKey, api_format: str, data: Dict[str, Any]) -> None:
|
||||
def _set_circuit_data(cls, key: ProviderAPIKey, api_format: str, data: dict[str, Any]) -> None:
|
||||
"""设置指定格式的熔断器数据"""
|
||||
circuit_by_format = dict(key.circuit_breaker_by_format or {})
|
||||
circuit_by_format[api_format] = data
|
||||
@@ -160,9 +160,9 @@ class HealthMonitor:
|
||||
def record_success(
|
||||
cls,
|
||||
db: Session,
|
||||
key_id: Optional[str] = None,
|
||||
api_format: Optional[str] = None,
|
||||
response_time_ms: Optional[int] = None,
|
||||
key_id: str | None = None,
|
||||
api_format: str | None = None,
|
||||
response_time_ms: int | None = None,
|
||||
) -> None:
|
||||
"""记录成功请求(按 API 格式)
|
||||
|
||||
@@ -285,9 +285,9 @@ class HealthMonitor:
|
||||
def record_failure(
|
||||
cls,
|
||||
db: Session,
|
||||
key_id: Optional[str] = None,
|
||||
api_format: Optional[str] = None,
|
||||
error_type: Optional[str] = None,
|
||||
key_id: str | None = None,
|
||||
api_format: str | None = None,
|
||||
error_type: str | None = None,
|
||||
) -> None:
|
||||
"""记录失败请求(按 API 格式)
|
||||
|
||||
@@ -433,7 +433,7 @@ class HealthMonitor:
|
||||
|
||||
@classmethod
|
||||
def _calculate_error_rate_from_window(
|
||||
cls, window: List[Dict[str, Any]], now_ts: float
|
||||
cls, window: list[dict[str, Any]], now_ts: float
|
||||
) -> float:
|
||||
"""从窗口数据计算错误率"""
|
||||
if not window:
|
||||
@@ -451,7 +451,7 @@ class HealthMonitor:
|
||||
# ==================== 熔断器状态方法(操作数据字典)====================
|
||||
|
||||
@classmethod
|
||||
def _get_circuit_state_from_data(cls, circuit_data: Dict[str, Any], now: datetime) -> str:
|
||||
def _get_circuit_state_from_data(cls, circuit_data: dict[str, Any], now: datetime) -> str:
|
||||
"""从数据字典获取当前熔断器状态"""
|
||||
if not circuit_data.get("open"):
|
||||
return CircuitState.CLOSED
|
||||
@@ -475,7 +475,7 @@ class HealthMonitor:
|
||||
@classmethod
|
||||
def _open_circuit_data(
|
||||
cls,
|
||||
circuit_data: Dict[str, Any],
|
||||
circuit_data: dict[str, Any],
|
||||
now: datetime,
|
||||
recovery_seconds: int,
|
||||
reason: str,
|
||||
@@ -489,7 +489,7 @@ class HealthMonitor:
|
||||
circuit_data["next_probe_at"] = (now + timedelta(seconds=recovery_seconds)).isoformat()
|
||||
|
||||
@classmethod
|
||||
def _enter_half_open_data(cls, circuit_data: Dict[str, Any], now: datetime) -> None:
|
||||
def _enter_half_open_data(cls, circuit_data: dict[str, Any], now: datetime) -> None:
|
||||
"""进入半开状态(操作数据字典)"""
|
||||
circuit_data["half_open_until"] = (
|
||||
now + timedelta(seconds=cls.HALF_OPEN_DURATION)
|
||||
@@ -499,7 +499,7 @@ class HealthMonitor:
|
||||
|
||||
@classmethod
|
||||
def _close_circuit_data(
|
||||
cls, circuit_data: Dict[str, Any], health_data: Dict[str, Any], reason: str
|
||||
cls, circuit_data: dict[str, Any], health_data: dict[str, Any], reason: str
|
||||
) -> None:
|
||||
"""关闭熔断器(操作数据字典)"""
|
||||
circuit_data["open"] = False
|
||||
@@ -527,7 +527,7 @@ class HealthMonitor:
|
||||
|
||||
@classmethod
|
||||
def is_circuit_breaker_closed(
|
||||
cls, resource: ProviderAPIKey, api_format: Optional[str] = None
|
||||
cls, resource: ProviderAPIKey, api_format: str | None = None
|
||||
) -> bool:
|
||||
"""检查熔断器是否允许请求通过(按 API 格式)"""
|
||||
if not api_format:
|
||||
@@ -564,8 +564,8 @@ class HealthMonitor:
|
||||
|
||||
@classmethod
|
||||
def get_circuit_breaker_status(
|
||||
cls, resource: ProviderAPIKey, api_format: Optional[str] = None
|
||||
) -> Tuple[bool, Optional[str]]:
|
||||
cls, resource: ProviderAPIKey, api_format: str | None = None
|
||||
) -> tuple[bool, str | None]:
|
||||
"""获取熔断器详细状态(按 API 格式)"""
|
||||
if not api_format:
|
||||
# 兼容旧调用:返回第一个开启的熔断器状态
|
||||
@@ -580,8 +580,8 @@ class HealthMonitor:
|
||||
|
||||
@classmethod
|
||||
def _get_status_from_circuit_data(
|
||||
cls, circuit_data: Dict[str, Any]
|
||||
) -> Tuple[bool, Optional[str]]:
|
||||
cls, circuit_data: dict[str, Any]
|
||||
) -> tuple[bool, str | None]:
|
||||
"""从熔断器数据获取状态描述"""
|
||||
if not circuit_data.get("open"):
|
||||
return True, None
|
||||
@@ -611,8 +611,8 @@ class HealthMonitor:
|
||||
|
||||
@classmethod
|
||||
def get_key_health(
|
||||
cls, db: Session, key_id: str, api_format: Optional[str] = None
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
cls, db: Session, key_id: str, api_format: str | None = None
|
||||
) -> dict[str, Any] | None:
|
||||
"""获取 Key 健康状态(支持按格式查询)"""
|
||||
try:
|
||||
key = db.query(ProviderAPIKey).filter(ProviderAPIKey.id == key_id).first()
|
||||
@@ -726,7 +726,7 @@ class HealthMonitor:
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def get_endpoint_health(cls, db: Session, endpoint_id: str) -> Optional[Dict[str, Any]]:
|
||||
def get_endpoint_health(cls, db: Session, endpoint_id: str) -> dict[str, Any] | None:
|
||||
"""获取 Endpoint 健康状态"""
|
||||
try:
|
||||
endpoint = (
|
||||
@@ -753,7 +753,7 @@ class HealthMonitor:
|
||||
|
||||
@classmethod
|
||||
def reset_health(
|
||||
cls, db: Session, key_id: Optional[str] = None, api_format: Optional[str] = None
|
||||
cls, db: Session, key_id: str | None = None, api_format: str | None = None
|
||||
) -> bool:
|
||||
"""重置健康度(支持按格式重置)"""
|
||||
try:
|
||||
@@ -781,7 +781,7 @@ class HealthMonitor:
|
||||
return False
|
||||
|
||||
@classmethod
|
||||
def manually_enable(cls, db: Session, key_id: Optional[str] = None) -> bool:
|
||||
def manually_enable(cls, db: Session, key_id: str | None = None) -> bool:
|
||||
"""手动启用 Key"""
|
||||
try:
|
||||
if key_id:
|
||||
@@ -803,7 +803,7 @@ class HealthMonitor:
|
||||
return False
|
||||
|
||||
@classmethod
|
||||
def get_all_health_status(cls, db: Session) -> Dict[str, Any]:
|
||||
def get_all_health_status(cls, db: Session) -> dict[str, Any]:
|
||||
"""获取所有健康状态摘要"""
|
||||
try:
|
||||
endpoint_stats = db.query(
|
||||
@@ -861,13 +861,13 @@ class HealthMonitor:
|
||||
# ==================== 历史记录方法 ====================
|
||||
|
||||
@classmethod
|
||||
def _push_circuit_event(cls, event: Dict[str, Any]) -> None:
|
||||
def _push_circuit_event(cls, event: dict[str, Any]) -> None:
|
||||
cls._circuit_history.append(event)
|
||||
if len(cls._circuit_history) > cls.CIRCUIT_HISTORY_LIMIT:
|
||||
cls._circuit_history.pop(0)
|
||||
|
||||
@classmethod
|
||||
def get_circuit_history(cls, limit: int = 50) -> List[Dict[str, Any]]:
|
||||
def get_circuit_history(cls, limit: int = 50) -> list[dict[str, Any]]:
|
||||
if limit <= 0:
|
||||
return []
|
||||
return cls._circuit_history[-limit:]
|
||||
@@ -878,9 +878,9 @@ class HealthMonitor:
|
||||
def is_eligible_for_probe(
|
||||
cls,
|
||||
db: Session,
|
||||
endpoint_id: Optional[str] = None,
|
||||
key_id: Optional[str] = None,
|
||||
api_format: Optional[str] = None,
|
||||
endpoint_id: str | None = None,
|
||||
key_id: str | None = None,
|
||||
api_format: str | None = None,
|
||||
) -> bool:
|
||||
"""检查是否有资格进行探测(按 API 格式)"""
|
||||
if not cls.ALLOW_AUTO_RECOVER:
|
||||
@@ -914,7 +914,7 @@ class HealthMonitor:
|
||||
|
||||
@classmethod
|
||||
def get_health_score(
|
||||
cls, key: ProviderAPIKey, api_format: Optional[str] = None
|
||||
cls, key: ProviderAPIKey, api_format: str | None = None
|
||||
) -> float:
|
||||
"""获取指定格式的健康度分数"""
|
||||
if not api_format:
|
||||
|
||||
@@ -2,7 +2,6 @@
|
||||
|
||||
import ipaddress
|
||||
from datetime import datetime, timezone
|
||||
from typing import Optional
|
||||
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
from sqlalchemy.orm import Session
|
||||
@@ -12,7 +11,7 @@ from src.core.logger import logger
|
||||
from src.models.database import ManagementToken
|
||||
|
||||
|
||||
def validate_ip_list(ips: Optional[list[str]]) -> Optional[list[str]]:
|
||||
def validate_ip_list(ips: list[str] | None) -> list[str] | None:
|
||||
"""验证 IP 白名单格式
|
||||
|
||||
- None: 不限制 IP
|
||||
@@ -45,7 +44,7 @@ def validate_ip_list(ips: Optional[list[str]]) -> Optional[list[str]]:
|
||||
return validated
|
||||
|
||||
|
||||
def parse_expires_at(v, allow_past: bool = False) -> Optional[datetime]:
|
||||
def parse_expires_at(v, allow_past: bool = False) -> datetime | None:
|
||||
"""解析过期时间,确保时区安全
|
||||
|
||||
前端 datetime-local 输入返回本地时间字符串(无时区信息)。
|
||||
@@ -88,7 +87,7 @@ def parse_expires_at(v, allow_past: bool = False) -> Optional[datetime]:
|
||||
|
||||
def token_to_dict(
|
||||
token: ManagementToken,
|
||||
raw_token: Optional[str] = None,
|
||||
raw_token: str | None = None,
|
||||
include_user: bool = False,
|
||||
) -> dict:
|
||||
"""将 ManagementToken 转换为字典
|
||||
@@ -136,9 +135,9 @@ class ManagementTokenService:
|
||||
db: Session,
|
||||
user_id: str,
|
||||
name: str,
|
||||
description: Optional[str] = None,
|
||||
allowed_ips: Optional[list[str]] = None,
|
||||
expires_at: Optional[datetime] = None,
|
||||
description: str | None = None,
|
||||
allowed_ips: list[str] | None = None,
|
||||
expires_at: datetime | None = None,
|
||||
) -> tuple[ManagementToken, str]:
|
||||
"""创建 Management Token
|
||||
|
||||
@@ -201,8 +200,8 @@ class ManagementTokenService:
|
||||
|
||||
@staticmethod
|
||||
def get_token_by_id(
|
||||
db: Session, token_id: str, user_id: Optional[str] = None
|
||||
) -> Optional[ManagementToken]:
|
||||
db: Session, token_id: str, user_id: str | None = None
|
||||
) -> ManagementToken | None:
|
||||
"""根据 ID 获取 Token
|
||||
|
||||
Args:
|
||||
@@ -221,8 +220,8 @@ class ManagementTokenService:
|
||||
@staticmethod
|
||||
def list_tokens(
|
||||
db: Session,
|
||||
user_id: Optional[str] = None,
|
||||
is_active: Optional[bool] = None,
|
||||
user_id: str | None = None,
|
||||
is_active: bool | None = None,
|
||||
skip: int = 0,
|
||||
limit: int = 100,
|
||||
) -> tuple[list[ManagementToken], int]:
|
||||
@@ -254,16 +253,16 @@ class ManagementTokenService:
|
||||
def update_token(
|
||||
db: Session,
|
||||
token_id: str,
|
||||
user_id: Optional[str] = None,
|
||||
name: Optional[str] = None,
|
||||
description: Optional[str] = None,
|
||||
allowed_ips: Optional[list[str]] = None,
|
||||
expires_at: Optional[datetime] = None,
|
||||
is_active: Optional[bool] = None,
|
||||
user_id: str | None = None,
|
||||
name: str | None = None,
|
||||
description: str | None = None,
|
||||
allowed_ips: list[str] | None = None,
|
||||
expires_at: datetime | None = None,
|
||||
is_active: bool | None = None,
|
||||
clear_description: bool = False,
|
||||
clear_allowed_ips: bool = False,
|
||||
clear_expires_at: bool = False,
|
||||
) -> Optional[ManagementToken]:
|
||||
) -> ManagementToken | None:
|
||||
"""更新 Token
|
||||
|
||||
Args:
|
||||
@@ -334,7 +333,7 @@ class ManagementTokenService:
|
||||
|
||||
@staticmethod
|
||||
def delete_token(
|
||||
db: Session, token_id: str, user_id: Optional[str] = None
|
||||
db: Session, token_id: str, user_id: str | None = None
|
||||
) -> bool:
|
||||
"""删除 Token
|
||||
|
||||
@@ -359,8 +358,8 @@ class ManagementTokenService:
|
||||
|
||||
@staticmethod
|
||||
def toggle_status(
|
||||
db: Session, token_id: str, user_id: Optional[str] = None
|
||||
) -> Optional[ManagementToken]:
|
||||
db: Session, token_id: str, user_id: str | None = None
|
||||
) -> ManagementToken | None:
|
||||
"""切换 Token 状态
|
||||
|
||||
Args:
|
||||
@@ -385,8 +384,8 @@ class ManagementTokenService:
|
||||
|
||||
@staticmethod
|
||||
def regenerate_token(
|
||||
db: Session, token_id: str, user_id: Optional[str] = None
|
||||
) -> tuple[Optional[ManagementToken], Optional[str], Optional[str]]:
|
||||
db: Session, token_id: str, user_id: str | None = None
|
||||
) -> tuple[ManagementToken | None, str | None, str | None]:
|
||||
"""重新生成 Token
|
||||
|
||||
Args:
|
||||
|
||||
@@ -13,7 +13,7 @@ Thinking 整流器(Rectifier)
|
||||
"""
|
||||
|
||||
import copy
|
||||
from typing import Any, Dict, List, Tuple
|
||||
from typing import Any
|
||||
|
||||
from src.core.logger import logger
|
||||
|
||||
@@ -27,7 +27,7 @@ class ThinkingRectifier:
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def rectify(request_body: Dict[str, Any]) -> Tuple[Dict[str, Any], bool]:
|
||||
def rectify(request_body: dict[str, Any]) -> tuple[dict[str, Any], bool]:
|
||||
"""
|
||||
整流请求体
|
||||
|
||||
@@ -68,7 +68,7 @@ class ThinkingRectifier:
|
||||
return rectified_body, modified
|
||||
|
||||
@staticmethod
|
||||
def _rectify_messages(messages: List[Dict[str, Any]]) -> Tuple[List[Dict[str, Any]], bool]:
|
||||
def _rectify_messages(messages: list[dict[str, Any]]) -> tuple[list[dict[str, Any]], bool]:
|
||||
"""
|
||||
整流消息列表
|
||||
|
||||
@@ -84,7 +84,7 @@ class ThinkingRectifier:
|
||||
return messages, False
|
||||
|
||||
modified = False
|
||||
result_messages: List[Dict[str, Any]] = []
|
||||
result_messages: list[dict[str, Any]] = []
|
||||
thinking_removed = 0
|
||||
signature_removed = 0
|
||||
|
||||
@@ -151,7 +151,7 @@ class ThinkingRectifier:
|
||||
return result_messages, modified
|
||||
|
||||
@staticmethod
|
||||
def _should_remove_top_level_thinking(body: Dict[str, Any]) -> bool:
|
||||
def _should_remove_top_level_thinking(body: dict[str, Any]) -> bool:
|
||||
"""
|
||||
判断是否应该删除顶层 thinking 参数
|
||||
|
||||
|
||||
@@ -8,9 +8,7 @@
|
||||
- API Key/User 的请求级访问限制由 models_service.AccessRestrictions 处理
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Optional
|
||||
|
||||
from sqlalchemy import or_
|
||||
from sqlalchemy.orm import Query, Session, contains_eager
|
||||
@@ -167,7 +165,7 @@ class ModelAvailabilityQuery:
|
||||
provider_ids: set[str],
|
||||
api_formats: list[str],
|
||||
provider_to_endpoint_formats: dict[str, set[str]],
|
||||
) -> dict[str, list[tuple[Optional[list[str]], set[str]]]]:
|
||||
) -> dict[str, list[tuple[list[str] | None, set[str]]]]:
|
||||
"""
|
||||
获取每个 Provider 的 Key 权限规则
|
||||
|
||||
@@ -197,7 +195,7 @@ class ModelAvailabilityQuery:
|
||||
.all()
|
||||
)
|
||||
|
||||
provider_key_rules: dict[str, list[tuple[Optional[list[str]], set[str]]]] = {}
|
||||
provider_key_rules: dict[str, list[tuple[list[str] | None, set[str]]]] = {}
|
||||
for key_id, provider_id, allowed_models_raw, key_formats in key_rows:
|
||||
if not provider_id:
|
||||
continue
|
||||
@@ -221,7 +219,7 @@ class ModelAvailabilityQuery:
|
||||
continue
|
||||
|
||||
# 类型兜底:allowed_models(安全优先)
|
||||
allowed_models: Optional[list[str]]
|
||||
allowed_models: list[str] | None
|
||||
if allowed_models_raw is None:
|
||||
# None = 不限制
|
||||
allowed_models = None
|
||||
@@ -242,7 +240,7 @@ class ModelAvailabilityQuery:
|
||||
def find_by_global_model_name(
|
||||
db: Session,
|
||||
model_name: str,
|
||||
provider_ids: Optional[set[str]] = None,
|
||||
provider_ids: set[str] | None = None,
|
||||
eager_load: bool = False,
|
||||
) -> Query:
|
||||
"""
|
||||
|
||||
@@ -8,8 +8,9 @@
|
||||
- 通过 PricingStrategy 抽象,支持自定义总输入上下文计算、缓存 TTL 差异化等
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Dict, Optional, Tuple, Union
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
@@ -17,7 +18,7 @@ from src.core.logger import logger
|
||||
from src.models.database import GlobalModel, Model, Provider
|
||||
|
||||
|
||||
ProviderRef = Union[str, Provider, None]
|
||||
ProviderRef = str | Provider | None
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -25,8 +26,8 @@ class TieredPriceResult:
|
||||
"""阶梯计费价格查询结果"""
|
||||
input_price_per_1m: float
|
||||
output_price_per_1m: float
|
||||
cache_creation_price_per_1m: Optional[float] = None
|
||||
cache_read_price_per_1m: Optional[float] = None
|
||||
cache_creation_price_per_1m: float | None = None
|
||||
cache_read_price_per_1m: float | None = None
|
||||
tier_index: int = 0 # 命中的阶梯索引
|
||||
|
||||
|
||||
@@ -45,9 +46,9 @@ class CostBreakdown:
|
||||
class ModelCostService:
|
||||
"""集中负责模型价格与成本计算,避免在 mapper/usage 中重复实现。"""
|
||||
|
||||
_price_cache: Dict[str, Dict[str, float]] = {}
|
||||
_cache_price_cache: Dict[str, Dict[str, float]] = {}
|
||||
_tiered_pricing_cache: Dict[str, Optional[dict]] = {}
|
||||
_price_cache: dict[str, dict[str, float]] = {}
|
||||
_cache_price_cache: dict[str, dict[str, float]] = {}
|
||||
_tiered_pricing_cache: dict[str, dict | None] = {}
|
||||
|
||||
def __init__(self, db: Session):
|
||||
self.db = db
|
||||
@@ -60,7 +61,7 @@ class ModelCostService:
|
||||
def get_tier_for_tokens(
|
||||
tiered_pricing: dict,
|
||||
total_input_tokens: int
|
||||
) -> Optional[dict]:
|
||||
) -> dict | None:
|
||||
"""
|
||||
根据总输入 token 数确定价格阶梯。
|
||||
|
||||
@@ -89,8 +90,8 @@ class ModelCostService:
|
||||
@staticmethod
|
||||
def get_cache_read_price_for_ttl(
|
||||
tier: dict,
|
||||
cache_ttl_minutes: Optional[int] = None
|
||||
) -> Optional[float]:
|
||||
cache_ttl_minutes: int | None = None
|
||||
) -> float | None:
|
||||
"""
|
||||
根据缓存 TTL 获取缓存读取价格。
|
||||
|
||||
@@ -122,7 +123,7 @@ class ModelCostService:
|
||||
|
||||
async def get_tiered_pricing_async(
|
||||
self, provider: ProviderRef, model: str
|
||||
) -> Optional[dict]:
|
||||
) -> dict | None:
|
||||
"""
|
||||
异步获取模型的阶梯计费配置。
|
||||
|
||||
@@ -138,7 +139,7 @@ class ModelCostService:
|
||||
|
||||
async def get_tiered_pricing_with_source_async(
|
||||
self, provider: ProviderRef, model: str
|
||||
) -> Optional[dict]:
|
||||
) -> dict | None:
|
||||
"""
|
||||
异步获取模型的阶梯计费配置及来源信息。
|
||||
|
||||
@@ -205,7 +206,7 @@ class ModelCostService:
|
||||
self._tiered_pricing_cache[cache_key] = result
|
||||
return result
|
||||
|
||||
def get_tiered_pricing(self, provider: ProviderRef, model: str) -> Optional[dict]:
|
||||
def get_tiered_pricing(self, provider: ProviderRef, model: str) -> dict | None:
|
||||
"""同步获取模型的阶梯计费配置(直接查缓存和数据库,避免事件循环开销)。
|
||||
|
||||
Returns:
|
||||
@@ -261,7 +262,7 @@ class ModelCostService:
|
||||
# 公共方法
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def get_model_price_async(self, provider: ProviderRef, model: str) -> Tuple[float, float]:
|
||||
async def get_model_price_async(self, provider: ProviderRef, model: str) -> tuple[float, float]:
|
||||
"""
|
||||
异步版本: 返回给定 provider/model 的 (input_price, output_price)。
|
||||
|
||||
@@ -352,7 +353,7 @@ class ModelCostService:
|
||||
self._price_cache[cache_key] = {"input": input_price, "output": output_price}
|
||||
return input_price, output_price
|
||||
|
||||
def get_model_price(self, provider: ProviderRef, model: str) -> Tuple[float, float]:
|
||||
def get_model_price(self, provider: ProviderRef, model: str) -> tuple[float, float]:
|
||||
"""
|
||||
返回给定 provider/model 的 (input_price, output_price)。
|
||||
直接查缓存和数据库,避免事件循环开销。
|
||||
@@ -435,7 +436,7 @@ class ModelCostService:
|
||||
|
||||
async def get_cache_prices_async(
|
||||
self, provider: ProviderRef, model: str, input_price: float
|
||||
) -> Tuple[Optional[float], Optional[float]]:
|
||||
) -> tuple[float | None, float | None]:
|
||||
"""
|
||||
异步版本: 返回缓存创建/读取价格(每 1M tokens)。
|
||||
|
||||
@@ -517,7 +518,7 @@ class ModelCostService:
|
||||
}
|
||||
return cache_creation_price, cache_read_price
|
||||
|
||||
async def get_request_price_async(self, provider: ProviderRef, model: str) -> Optional[float]:
|
||||
async def get_request_price_async(self, provider: ProviderRef, model: str) -> float | None:
|
||||
"""
|
||||
异步版本: 返回按次计费价格(每次请求的固定费用)。
|
||||
|
||||
@@ -563,7 +564,7 @@ class ModelCostService:
|
||||
|
||||
return price_per_request
|
||||
|
||||
def get_request_price(self, provider: ProviderRef, model: str) -> Optional[float]:
|
||||
def get_request_price(self, provider: ProviderRef, model: str) -> float | None:
|
||||
"""
|
||||
返回按次计费价格(每次请求的固定费用)。
|
||||
直接查数据库,避免事件循环开销。
|
||||
@@ -608,7 +609,7 @@ class ModelCostService:
|
||||
|
||||
def get_cache_prices(
|
||||
self, provider: ProviderRef, model: str, input_price: float
|
||||
) -> Tuple[Optional[float], Optional[float]]:
|
||||
) -> tuple[float | None, float | None]:
|
||||
"""
|
||||
返回缓存创建/读取价格(每 1M tokens)。
|
||||
直接查缓存和数据库,避免事件循环开销。
|
||||
@@ -689,7 +690,7 @@ class ModelCostService:
|
||||
model: str,
|
||||
input_tokens: int,
|
||||
output_tokens: int,
|
||||
) -> Dict[str, float]:
|
||||
) -> dict[str, float]:
|
||||
"""返回与旧 ModelMapper.calculate_cost 相同结构的费用信息。"""
|
||||
input_price, output_price = self.get_model_price(provider, model)
|
||||
input_cost, output_cost, _, _, _, _, total_cost = self.compute_cost(
|
||||
@@ -715,10 +716,10 @@ class ModelCostService:
|
||||
output_price_per_1m: float,
|
||||
cache_creation_input_tokens: int = 0,
|
||||
cache_read_input_tokens: int = 0,
|
||||
cache_creation_price_per_1m: Optional[float] = None,
|
||||
cache_read_price_per_1m: Optional[float] = None,
|
||||
price_per_request: Optional[float] = None,
|
||||
) -> Tuple[float, float, float, float, float, float, float]:
|
||||
cache_creation_price_per_1m: float | None = None,
|
||||
cache_read_price_per_1m: float | None = None,
|
||||
price_per_request: float | None = None,
|
||||
) -> tuple[float, float, float, float, float, float, float]:
|
||||
"""成本计算核心逻辑(固定价格模式),供 UsageService 等复用。
|
||||
|
||||
Returns:
|
||||
@@ -760,15 +761,15 @@ class ModelCostService:
|
||||
output_tokens: int,
|
||||
cache_creation_input_tokens: int = 0,
|
||||
cache_read_input_tokens: int = 0,
|
||||
tiered_pricing: Optional[dict] = None,
|
||||
cache_ttl_minutes: Optional[int] = None,
|
||||
price_per_request: Optional[float] = None,
|
||||
tiered_pricing: dict | None = None,
|
||||
cache_ttl_minutes: int | None = None,
|
||||
price_per_request: float | None = None,
|
||||
# 回退价格(当没有阶梯配置时使用)
|
||||
fallback_input_price_per_1m: float = 0.0,
|
||||
fallback_output_price_per_1m: float = 0.0,
|
||||
fallback_cache_creation_price_per_1m: Optional[float] = None,
|
||||
fallback_cache_read_price_per_1m: Optional[float] = None,
|
||||
) -> Tuple[float, float, float, float, float, float, float, Optional[int]]:
|
||||
fallback_cache_creation_price_per_1m: float | None = None,
|
||||
fallback_cache_read_price_per_1m: float | None = None,
|
||||
) -> tuple[float, float, float, float, float, float, float, int | None]:
|
||||
"""
|
||||
支持阶梯计费的成本计算核心逻辑。
|
||||
|
||||
@@ -872,7 +873,7 @@ class ModelCostService:
|
||||
return provider.name
|
||||
return provider or "unknown"
|
||||
|
||||
def _resolve_provider(self, provider: ProviderRef) -> Optional[Provider]:
|
||||
def _resolve_provider(self, provider: ProviderRef) -> Provider | None:
|
||||
if isinstance(provider, Provider):
|
||||
return provider
|
||||
if not provider or provider == "unknown":
|
||||
@@ -891,9 +892,9 @@ class ModelCostService:
|
||||
output_tokens: int,
|
||||
cache_creation_input_tokens: int = 0,
|
||||
cache_read_input_tokens: int = 0,
|
||||
api_format: Optional[str] = None,
|
||||
cache_ttl_minutes: Optional[int] = None,
|
||||
) -> Tuple[float, float, float, float, float, float, float, Optional[int]]:
|
||||
api_format: str | None = None,
|
||||
cache_ttl_minutes: int | None = None,
|
||||
) -> tuple[float, float, float, float, float, float, float, int | None]:
|
||||
"""
|
||||
使用计费策略计算成本(异步版本)
|
||||
|
||||
@@ -981,21 +982,15 @@ class ModelCostService:
|
||||
output_tokens: int,
|
||||
cache_creation_input_tokens: int = 0,
|
||||
cache_read_input_tokens: int = 0,
|
||||
api_format: Optional[str] = None,
|
||||
cache_ttl_minutes: Optional[int] = None,
|
||||
) -> Tuple[float, float, float, float, float, float, float, Optional[int]]:
|
||||
api_format: str | None = None,
|
||||
cache_ttl_minutes: int | None = None,
|
||||
) -> tuple[float, float, float, float, float, float, float, int | None]:
|
||||
"""
|
||||
使用计费策略计算成本(同步版本)
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
try:
|
||||
loop = asyncio.get_event_loop()
|
||||
except RuntimeError:
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
|
||||
return loop.run_until_complete(
|
||||
return asyncio.run(
|
||||
self.compute_cost_with_strategy_async(
|
||||
provider=provider,
|
||||
model=model,
|
||||
|
||||
@@ -15,7 +15,6 @@ import asyncio
|
||||
import fnmatch
|
||||
import os
|
||||
from datetime import datetime, timezone
|
||||
from typing import List, Optional, Set
|
||||
|
||||
from sqlalchemy.orm import Session, joinedload
|
||||
|
||||
@@ -63,10 +62,10 @@ def _match_pattern(model_id: str, pattern: str) -> bool:
|
||||
|
||||
|
||||
def _filter_models_by_patterns(
|
||||
model_ids: Set[str],
|
||||
include_patterns: Optional[List[str]],
|
||||
exclude_patterns: Optional[List[str]],
|
||||
) -> Set[str]:
|
||||
model_ids: set[str],
|
||||
include_patterns: list[str] | None,
|
||||
exclude_patterns: list[str] | None,
|
||||
) -> set[str]:
|
||||
"""
|
||||
根据包含/排除规则过滤模型列表
|
||||
|
||||
@@ -113,7 +112,7 @@ def _get_upstream_models_cache_key(provider_id: str, api_key_id: str) -> str:
|
||||
|
||||
async def get_upstream_models_from_cache(
|
||||
provider_id: str, api_key_id: str
|
||||
) -> Optional[list[dict]]:
|
||||
) -> list[dict] | None:
|
||||
"""从缓存获取上游模型列表"""
|
||||
cache_key = _get_upstream_models_cache_key(provider_id, api_key_id)
|
||||
cached = await CacheService.get(cache_key)
|
||||
@@ -138,7 +137,7 @@ class ModelFetchScheduler:
|
||||
def __init__(self) -> None:
|
||||
self._running = False
|
||||
self._lock = asyncio.Lock()
|
||||
self._startup_task: Optional[asyncio.Task] = None
|
||||
self._startup_task: asyncio.Task | None = None
|
||||
|
||||
async def start(self) -> None:
|
||||
"""启动调度器"""
|
||||
@@ -238,7 +237,7 @@ class ModelFetchScheduler:
|
||||
skip_count += 1
|
||||
else:
|
||||
error_count += 1
|
||||
except asyncio.TimeoutError:
|
||||
except TimeoutError:
|
||||
logger.error(f"处理 Key {key_id} 超时({KEY_FETCH_TIMEOUT_SECONDS}s)")
|
||||
self._update_key_error(key_id, f"Timeout after {KEY_FETCH_TIMEOUT_SECONDS}s")
|
||||
error_count += 1
|
||||
@@ -290,7 +289,7 @@ class ModelFetchScheduler:
|
||||
|
||||
async def _fetch_models_for_key(
|
||||
self,
|
||||
db: "Session",
|
||||
db: Session,
|
||||
key: ProviderAPIKey,
|
||||
) -> str:
|
||||
"""为单个 Key 获取模型并更新 allowed_models,返回结果状态"""
|
||||
@@ -455,7 +454,7 @@ class ModelFetchScheduler:
|
||||
|
||||
|
||||
# 单例模式
|
||||
_model_fetch_scheduler: Optional[ModelFetchScheduler] = None
|
||||
_model_fetch_scheduler: ModelFetchScheduler | None = None
|
||||
|
||||
|
||||
def get_model_fetch_scheduler() -> ModelFetchScheduler:
|
||||
|
||||
@@ -4,9 +4,8 @@ GlobalModel 服务层
|
||||
提供 GlobalModel 的 CRUD 操作、查询和统计功能
|
||||
"""
|
||||
|
||||
from typing import Dict, List, Optional, Set, cast
|
||||
from typing import cast
|
||||
|
||||
from sqlalchemy import and_, func
|
||||
from sqlalchemy.orm import Session, joinedload
|
||||
|
||||
from src.core.exceptions import InvalidRequestException, NotFoundException
|
||||
@@ -18,7 +17,7 @@ from src.models.pydantic_models import GlobalModelUpdate
|
||||
async def on_key_allowed_models_changed(
|
||||
db: Session,
|
||||
provider_id: str,
|
||||
allowed_models: Optional[List[str]] = None,
|
||||
allowed_models: list[str] | None = None,
|
||||
skip_disassociate: bool = False,
|
||||
) -> None:
|
||||
"""
|
||||
@@ -84,7 +83,7 @@ class GlobalModelService:
|
||||
return global_model
|
||||
|
||||
@staticmethod
|
||||
def get_global_model_by_name(db: Session, name: str) -> Optional[GlobalModel]:
|
||||
def get_global_model_by_name(db: Session, name: str) -> GlobalModel | None:
|
||||
"""通过名称获取 GlobalModel"""
|
||||
return db.query(GlobalModel).filter(GlobalModel.name == name).first()
|
||||
|
||||
@@ -93,9 +92,9 @@ class GlobalModelService:
|
||||
db: Session,
|
||||
skip: int = 0,
|
||||
limit: int = 100,
|
||||
is_active: Optional[bool] = None,
|
||||
search: Optional[str] = None,
|
||||
) -> List[GlobalModel]:
|
||||
is_active: bool | None = None,
|
||||
search: str | None = None,
|
||||
) -> list[GlobalModel]:
|
||||
"""列出 GlobalModel"""
|
||||
query = db.query(GlobalModel)
|
||||
|
||||
@@ -119,15 +118,15 @@ class GlobalModelService:
|
||||
db: Session,
|
||||
name: str,
|
||||
display_name: str,
|
||||
is_active: Optional[bool] = True,
|
||||
is_active: bool | None = True,
|
||||
# 按次计费配置
|
||||
default_price_per_request: Optional[float] = None,
|
||||
default_price_per_request: float | None = None,
|
||||
# 阶梯计费配置(必填)
|
||||
default_tiered_pricing: dict = None,
|
||||
# Key 能力配置
|
||||
supported_capabilities: Optional[List[str]] = None,
|
||||
supported_capabilities: list[str] | None = None,
|
||||
# 模型配置(JSON)
|
||||
config: Optional[dict] = None,
|
||||
config: dict | None = None,
|
||||
) -> GlobalModel:
|
||||
"""创建 GlobalModel"""
|
||||
# 检查名称是否已存在
|
||||
@@ -219,7 +218,7 @@ class GlobalModelService:
|
||||
db.commit()
|
||||
|
||||
@staticmethod
|
||||
def get_global_model_stats(db: Session, global_model_id: str) -> Dict:
|
||||
def get_global_model_stats(db: Session, global_model_id: str) -> dict:
|
||||
"""获取 GlobalModel 统计信息"""
|
||||
global_model = GlobalModelService.get_global_model(db, global_model_id)
|
||||
|
||||
@@ -232,7 +231,7 @@ class GlobalModelService:
|
||||
)
|
||||
|
||||
# 统计支持的 Provider 数量
|
||||
provider_ids = set(model.provider_id for model in models)
|
||||
provider_ids = {model.provider_id for model in models}
|
||||
|
||||
# 从阶梯计费中提取价格范围
|
||||
input_prices = []
|
||||
@@ -263,11 +262,10 @@ class GlobalModelService:
|
||||
def batch_assign_to_providers(
|
||||
db: Session,
|
||||
global_model_id: str,
|
||||
provider_ids: List[str],
|
||||
provider_ids: list[str],
|
||||
create_models: bool = False,
|
||||
) -> Dict:
|
||||
) -> dict:
|
||||
"""批量为多个 Provider 添加 GlobalModel 实现"""
|
||||
from .service import ModelService
|
||||
|
||||
global_model = GlobalModelService.get_global_model(db, global_model_id)
|
||||
|
||||
@@ -338,8 +336,8 @@ class GlobalModelService:
|
||||
def auto_associate_provider_by_key_whitelist(
|
||||
db: Session,
|
||||
provider_id: str,
|
||||
allowed_models: List[str],
|
||||
) -> Dict:
|
||||
allowed_models: list[str],
|
||||
) -> dict:
|
||||
"""
|
||||
根据 Key 白名单自动关联 Provider 到匹配的 GlobalModel
|
||||
|
||||
@@ -358,7 +356,7 @@ class GlobalModelService:
|
||||
from src.core.model_permissions import match_model_with_pattern
|
||||
from src.models.database import Provider
|
||||
|
||||
results: Dict[str, List[Dict]] = {
|
||||
results: dict[str, list[dict]] = {
|
||||
"success": [],
|
||||
"errors": [],
|
||||
}
|
||||
@@ -378,9 +376,9 @@ class GlobalModelService:
|
||||
.filter(Model.provider_id == provider_id)
|
||||
.all()
|
||||
)
|
||||
linked_global_model_ids: Set[str] = {row[0] for row in existing_associations if row[0]}
|
||||
linked_global_model_ids: set[str] = {row[0] for row in existing_associations if row[0]}
|
||||
# 同时获取已存在的 provider_model_name 集合,避免唯一约束冲突
|
||||
existing_provider_model_names: Set[str] = {
|
||||
existing_provider_model_names: set[str] = {
|
||||
row[1] for row in existing_associations if row[1]
|
||||
}
|
||||
|
||||
@@ -403,7 +401,7 @@ class GlobalModelService:
|
||||
continue
|
||||
|
||||
# 提取映射规则
|
||||
model_mappings: List[str] = []
|
||||
model_mappings: list[str] = []
|
||||
if global_model.config and isinstance(global_model.config, dict):
|
||||
mappings = global_model.config.get("model_mappings")
|
||||
if isinstance(mappings, list):
|
||||
@@ -472,7 +470,7 @@ class GlobalModelService:
|
||||
def auto_disassociate_provider_by_key_whitelist(
|
||||
db: Session,
|
||||
provider_id: str,
|
||||
) -> Dict:
|
||||
) -> dict:
|
||||
"""
|
||||
根据 Key 白名单自动解除 Provider 与不再匹配的 GlobalModel 的关联
|
||||
|
||||
@@ -491,7 +489,7 @@ class GlobalModelService:
|
||||
from src.core.model_permissions import match_model_with_pattern
|
||||
from src.models.database import Provider, ProviderAPIKey
|
||||
|
||||
results: Dict[str, List[Dict]] = {
|
||||
results: dict[str, list[dict]] = {
|
||||
"success": [],
|
||||
"errors": [],
|
||||
}
|
||||
@@ -514,7 +512,7 @@ class GlobalModelService:
|
||||
|
||||
# 收集所有 Key 的 allowed_models 并集
|
||||
# 注意:allowed_models 为 null 表示允许所有模型,此时不应解除任何关联
|
||||
all_allowed_models: Set[str] = set()
|
||||
all_allowed_models: set[str] = set()
|
||||
has_unlimited_key = False # 是否存在允许所有模型的 Key
|
||||
|
||||
for key in keys:
|
||||
@@ -542,7 +540,7 @@ class GlobalModelService:
|
||||
)
|
||||
|
||||
# 3. 检查每个 Model 是否还能匹配,收集需要删除的 Model
|
||||
models_to_delete: List[Model] = []
|
||||
models_to_delete: list[Model] = []
|
||||
|
||||
for model in models:
|
||||
# 跳过没有关联 GlobalModel 的
|
||||
@@ -552,7 +550,7 @@ class GlobalModelService:
|
||||
global_model = cast(GlobalModel, model.global_model)
|
||||
|
||||
# 提取映射规则
|
||||
model_mappings: List[str] = []
|
||||
model_mappings: list[str] = []
|
||||
config = global_model.config
|
||||
if config and isinstance(config, dict):
|
||||
mappings = config.get("model_mappings")
|
||||
|
||||
@@ -3,8 +3,6 @@
|
||||
根据数据库中的配置,将用户请求的模型映射到提供商的实际模型
|
||||
"""
|
||||
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
from sqlalchemy.orm import Session, joinedload
|
||||
|
||||
from src.core.cache_utils import SyncLRUCache
|
||||
@@ -79,7 +77,7 @@ class ModelMapperMiddleware:
|
||||
|
||||
async def get_mapping(
|
||||
self, source_model: str, provider_id: str
|
||||
) -> Optional[object]:
|
||||
) -> object | None:
|
||||
"""
|
||||
获取模型映射
|
||||
|
||||
@@ -137,7 +135,7 @@ class ModelMapperMiddleware:
|
||||
|
||||
return mapping
|
||||
|
||||
def get_all_mappings(self, provider_id: str) -> List[object]:
|
||||
def get_all_mappings(self, provider_id: str) -> list[object]:
|
||||
"""
|
||||
获取提供商的所有可用模型(通过 GlobalModel)
|
||||
|
||||
@@ -177,7 +175,7 @@ class ModelMapperMiddleware:
|
||||
|
||||
return mappings
|
||||
|
||||
def get_supported_models(self, provider_id: str) -> List[str]:
|
||||
def get_supported_models(self, provider_id: str) -> list[str]:
|
||||
"""
|
||||
获取提供商支持的所有源模型名
|
||||
|
||||
@@ -192,7 +190,7 @@ class ModelMapperMiddleware:
|
||||
|
||||
async def validate_request(
|
||||
self, request: ClaudeMessagesRequest, provider: Provider
|
||||
) -> tuple[bool, Optional[str]]:
|
||||
) -> tuple[bool, str | None]:
|
||||
"""
|
||||
验证请求是否符合映射的限制
|
||||
|
||||
@@ -219,7 +217,7 @@ class ModelMapperMiddleware:
|
||||
self._cache.clear()
|
||||
logger.debug("Model mapping cache cleared")
|
||||
|
||||
def refresh_cache(self, provider_id: Optional[str] = None):
|
||||
def refresh_cache(self, provider_id: str | None = None):
|
||||
"""
|
||||
刷新缓存
|
||||
|
||||
@@ -258,10 +256,10 @@ class ModelRoutingMiddleware:
|
||||
def select_provider(
|
||||
self,
|
||||
model_name: str,
|
||||
preferred_provider: Optional[str] = None,
|
||||
allowed_api_formats: Optional[List[str]] = None,
|
||||
request_id: Optional[str] = None,
|
||||
) -> Optional[Provider]:
|
||||
preferred_provider: str | None = None,
|
||||
allowed_api_formats: list[str] | None = None,
|
||||
request_id: str | None = None,
|
||||
) -> Provider | None:
|
||||
"""
|
||||
根据模型名选择提供商
|
||||
|
||||
@@ -327,7 +325,7 @@ class ModelRoutingMiddleware:
|
||||
logger.error("No active providers found.")
|
||||
return None
|
||||
|
||||
def get_available_models(self) -> Dict[str, List[str]]:
|
||||
def get_available_models(self) -> dict[str, list[str]]:
|
||||
"""
|
||||
获取所有可用的模型及其提供商
|
||||
|
||||
@@ -354,7 +352,7 @@ class ModelRoutingMiddleware:
|
||||
|
||||
return result
|
||||
|
||||
async def get_cheapest_provider(self, model_name: str) -> Optional[Provider]:
|
||||
async def get_cheapest_provider(self, model_name: str) -> Provider | None:
|
||||
"""
|
||||
获取某个模型最便宜的提供商
|
||||
|
||||
|
||||
@@ -11,7 +11,6 @@
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -28,11 +27,11 @@ class PricingConfig:
|
||||
"""价格配置"""
|
||||
input_price_per_1m: float = 0.0
|
||||
output_price_per_1m: float = 0.0
|
||||
cache_creation_price_per_1m: Optional[float] = None
|
||||
cache_read_price_per_1m: Optional[float] = None
|
||||
price_per_request: Optional[float] = None
|
||||
tiered_pricing: Optional[dict] = None
|
||||
cache_ttl_minutes: Optional[int] = None
|
||||
cache_creation_price_per_1m: float | None = None
|
||||
cache_read_price_per_1m: float | None = None
|
||||
price_per_request: float | None = None
|
||||
tiered_pricing: dict | None = None
|
||||
cache_ttl_minutes: int | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -45,4 +44,4 @@ class CostResult:
|
||||
cache_cost: float = 0.0
|
||||
request_cost: float = 0.0
|
||||
total_cost: float = 0.0
|
||||
tier_index: Optional[int] = None # 命中的阶梯索引
|
||||
tier_index: int | None = None # 命中的阶梯索引
|
||||
|
||||
@@ -3,7 +3,6 @@
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from typing import List, Optional
|
||||
|
||||
from sqlalchemy import and_
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
@@ -126,8 +125,8 @@ class ModelService:
|
||||
provider_id: str, # UUID
|
||||
skip: int = 0,
|
||||
limit: int = 100,
|
||||
is_active: Optional[bool] = None,
|
||||
) -> List[Model]:
|
||||
is_active: bool | None = None,
|
||||
) -> list[Model]:
|
||||
"""获取提供商的模型列表"""
|
||||
from sqlalchemy.orm import joinedload
|
||||
|
||||
@@ -150,9 +149,9 @@ class ModelService:
|
||||
db: Session,
|
||||
skip: int = 0,
|
||||
limit: int = 100,
|
||||
is_active: Optional[bool] = None,
|
||||
category: Optional[str] = None,
|
||||
) -> List[Model]:
|
||||
is_active: bool | None = None,
|
||||
category: str | None = None,
|
||||
) -> list[Model]:
|
||||
"""获取所有模型列表"""
|
||||
query = db.query(Model)
|
||||
|
||||
@@ -336,7 +335,7 @@ class ModelService:
|
||||
return model
|
||||
|
||||
@staticmethod
|
||||
def get_model_by_name(db: Session, provider_id: str, model_name: str) -> Optional[Model]:
|
||||
def get_model_by_name(db: Session, provider_id: str, model_name: str) -> Model | None:
|
||||
"""根据 provider_model_name 获取模型"""
|
||||
return (
|
||||
db.query(Model)
|
||||
@@ -346,8 +345,8 @@ class ModelService:
|
||||
|
||||
@staticmethod
|
||||
def batch_create_models(
|
||||
db: Session, provider_id: str, models_data: List[ModelCreate]
|
||||
) -> List[Model]: # UUID
|
||||
db: Session, provider_id: str, models_data: list[ModelCreate]
|
||||
) -> list[Model]: # UUID
|
||||
"""批量创建模型"""
|
||||
# 检查提供商是否存在
|
||||
provider = db.query(Provider).filter(Provider.id == provider_id).first()
|
||||
|
||||
@@ -7,7 +7,6 @@
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from typing import Dict, Optional
|
||||
|
||||
import httpx
|
||||
|
||||
@@ -23,7 +22,7 @@ MAX_CONCURRENT_REQUESTS = 5
|
||||
MODEL_FETCH_FORMATS = [APIFormat.OPENAI, APIFormat.CLAUDE, APIFormat.GEMINI]
|
||||
|
||||
|
||||
def _get_adapter_for_format(api_format: str) -> Optional[type]:
|
||||
def _get_adapter_for_format(api_format: str) -> type | None:
|
||||
"""根据 API 格式获取对应的 Adapter 类"""
|
||||
from src.api.handlers.base.chat_adapter_base import get_adapter_class
|
||||
from src.api.handlers.base.cli_adapter_base import get_cli_adapter_class
|
||||
@@ -39,7 +38,7 @@ def _get_adapter_for_format(api_format: str) -> Optional[type]:
|
||||
|
||||
def build_all_format_configs(
|
||||
api_key_value: str,
|
||||
format_to_endpoint: Dict[str, ProviderEndpoint],
|
||||
format_to_endpoint: dict[str, ProviderEndpoint],
|
||||
) -> list[dict]:
|
||||
"""
|
||||
构建所有 API 格式的端点配置
|
||||
@@ -118,7 +117,7 @@ async def fetch_models_from_endpoints(
|
||||
|
||||
async def fetch_one(
|
||||
client: httpx.AsyncClient, config: dict
|
||||
) -> tuple[list, Optional[str], bool]:
|
||||
) -> tuple[list, str | None, bool]:
|
||||
base_url = config["base_url"]
|
||||
if not base_url:
|
||||
return [], None, False
|
||||
|
||||
@@ -6,7 +6,7 @@
|
||||
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
@@ -48,11 +48,11 @@ class CandidateResolver:
|
||||
api_format: APIFormat,
|
||||
model_name: str,
|
||||
affinity_key: str,
|
||||
user_api_key: Optional[ApiKey] = None,
|
||||
request_id: Optional[str] = None,
|
||||
user_api_key: ApiKey | None = None,
|
||||
request_id: str | None = None,
|
||||
is_stream: bool = False,
|
||||
capability_requirements: Optional[Dict[str, bool]] = None,
|
||||
) -> Tuple[List[ProviderCandidate], str]:
|
||||
capability_requirements: dict[str, bool] | None = None,
|
||||
) -> tuple[list[ProviderCandidate], str]:
|
||||
"""
|
||||
获取所有可用候选
|
||||
|
||||
@@ -71,10 +71,10 @@ class CandidateResolver:
|
||||
Raises:
|
||||
ProviderNotAvailableException: 没有找到任何可用候选时
|
||||
"""
|
||||
all_candidates: List[ProviderCandidate] = []
|
||||
all_candidates: list[ProviderCandidate] = []
|
||||
provider_offset = 0
|
||||
provider_batch_size = 20
|
||||
global_model_id: Optional[str] = None
|
||||
global_model_id: str | None = None
|
||||
|
||||
while True:
|
||||
candidates, resolved_global_model_id = await self.cache_scheduler.list_all_candidates(
|
||||
@@ -112,12 +112,12 @@ class CandidateResolver:
|
||||
|
||||
def create_candidate_records(
|
||||
self,
|
||||
all_candidates: List[ProviderCandidate],
|
||||
request_id: Optional[str],
|
||||
all_candidates: list[ProviderCandidate],
|
||||
request_id: str | None,
|
||||
user_id: str,
|
||||
user_api_key: ApiKey,
|
||||
required_capabilities: Optional[Dict[str, bool]] = None,
|
||||
) -> Dict[Tuple[int, int], str]:
|
||||
required_capabilities: dict[str, bool] | None = None,
|
||||
) -> dict[tuple[int, int], str]:
|
||||
"""
|
||||
为所有候选预先创建 available 状态记录(批量插入优化)
|
||||
|
||||
@@ -133,8 +133,8 @@ class CandidateResolver:
|
||||
"""
|
||||
from src.models.database import RequestCandidate
|
||||
|
||||
candidate_records_to_insert: List[Dict[str, Any]] = []
|
||||
candidate_record_map: Dict[Tuple[int, int], str] = {}
|
||||
candidate_records_to_insert: list[dict[str, Any]] = []
|
||||
candidate_record_map: dict[tuple[int, int], str] = {}
|
||||
|
||||
# 只保存启用的能力(值为 True 的)
|
||||
active_capabilities = None
|
||||
@@ -208,8 +208,8 @@ class CandidateResolver:
|
||||
|
||||
def get_active_candidates(
|
||||
self,
|
||||
all_candidates: List[ProviderCandidate],
|
||||
) -> List[Tuple[int, ProviderCandidate]]:
|
||||
all_candidates: list[ProviderCandidate],
|
||||
) -> list[tuple[int, ProviderCandidate]]:
|
||||
"""
|
||||
获取所有非跳过的候选(带索引)
|
||||
|
||||
@@ -223,7 +223,7 @@ class CandidateResolver:
|
||||
|
||||
def count_total_attempts(
|
||||
self,
|
||||
all_candidates: List[ProviderCandidate],
|
||||
all_candidates: list[ProviderCandidate],
|
||||
) -> int:
|
||||
"""
|
||||
计算总尝试次数
|
||||
|
||||
@@ -6,7 +6,7 @@
|
||||
|
||||
import json
|
||||
from enum import Enum
|
||||
from typing import Any, Dict, Optional, Tuple, Union
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
from sqlalchemy.orm import Session
|
||||
@@ -51,7 +51,7 @@ class ErrorClassifier:
|
||||
"""
|
||||
|
||||
# 需要触发故障转移的错误类型
|
||||
RETRIABLE_ERRORS: Tuple[type, ...] = (
|
||||
RETRIABLE_ERRORS: tuple[type, ...] = (
|
||||
ProviderException, # 包含所有 Provider 异常子类
|
||||
ConnectionError, # Python 标准连接错误
|
||||
TimeoutError, # Python 标准超时错误
|
||||
@@ -59,7 +59,7 @@ class ErrorClassifier:
|
||||
)
|
||||
|
||||
# 不可重试的错误类型(直接抛出)
|
||||
NON_RETRIABLE_ERRORS: Tuple[type, ...] = (
|
||||
NON_RETRIABLE_ERRORS: tuple[type, ...] = (
|
||||
ValueError, # 参数错误
|
||||
TypeError, # 类型错误
|
||||
KeyError, # 键错误
|
||||
@@ -73,7 +73,7 @@ class ErrorClassifier:
|
||||
#
|
||||
# 重要:不要在此列表中包含 Provider Key 配置问题(如 invalid_api_key)
|
||||
# 这类错误应该触发故障转移,而不是直接返回给用户
|
||||
CLIENT_ERROR_PATTERNS: Tuple[str, ...] = (
|
||||
CLIENT_ERROR_PATTERNS: tuple[str, ...] = (
|
||||
"could not process image", # 图片处理失败
|
||||
"image too large", # 图片过大
|
||||
"invalid image", # 无效图片
|
||||
@@ -101,7 +101,7 @@ class ErrorClassifier:
|
||||
self,
|
||||
db: Session,
|
||||
adaptive_manager: Any = None,
|
||||
cache_scheduler: Optional[CacheAwareScheduler] = None,
|
||||
cache_scheduler: CacheAwareScheduler | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
初始化错误分类器
|
||||
@@ -117,7 +117,7 @@ class ErrorClassifier:
|
||||
|
||||
# 表示客户端错误的 error type(不区分大小写)
|
||||
# 这些 type 表明是请求本身的问题,不应重试
|
||||
CLIENT_ERROR_TYPES: Tuple[str, ...] = (
|
||||
CLIENT_ERROR_TYPES: tuple[str, ...] = (
|
||||
# Claude/OpenAI 标准
|
||||
"invalid_request_error",
|
||||
# Gemini
|
||||
@@ -131,7 +131,7 @@ class ErrorClassifier:
|
||||
)
|
||||
|
||||
# 表示客户端错误的 reason/code 字段值
|
||||
CLIENT_ERROR_REASONS: Tuple[str, ...] = (
|
||||
CLIENT_ERROR_REASONS: tuple[str, ...] = (
|
||||
"CONTENT_LENGTH_EXCEEDS_THRESHOLD",
|
||||
"CONTEXT_LENGTH_EXCEEDED",
|
||||
"MAX_TOKENS_EXCEEDED",
|
||||
@@ -141,7 +141,7 @@ class ErrorClassifier:
|
||||
|
||||
# Provider 兼容性错误模式 - 这类错误应该触发故障转移
|
||||
# 因为换一个 Provider 可能就能成功
|
||||
COMPATIBILITY_ERROR_PATTERNS: Tuple[str, ...] = (
|
||||
COMPATIBILITY_ERROR_PATTERNS: tuple[str, ...] = (
|
||||
"unsupported parameter", # 不支持的参数
|
||||
"unsupported model", # 不支持的模型
|
||||
"unsupported feature", # 不支持的功能
|
||||
@@ -154,7 +154,7 @@ class ErrorClassifier:
|
||||
|
||||
# Thinking 块相关错误模式 - 这类错误需要清洗 thinking 块或调整请求
|
||||
# 场景:多供应商环境下,Provider A 生成的 thinking 块被发送到 Provider B 时签名验证失败
|
||||
THINKING_ERROR_PATTERNS: Tuple[str, ...] = (
|
||||
THINKING_ERROR_PATTERNS: tuple[str, ...] = (
|
||||
# 签名错误:跨 Provider 发送 thinking 块时,签名无法被目标 Provider 验证
|
||||
# 例: "invalid `signature` in `thinking` block: signature is for a different request"
|
||||
"invalid `signature` in `thinking` block",
|
||||
@@ -176,7 +176,7 @@ class ErrorClassifier:
|
||||
"expected `redacted_thinking`, found",
|
||||
)
|
||||
|
||||
def _parse_error_response(self, error_text: Optional[str]) -> Dict[str, Any]:
|
||||
def _parse_error_response(self, error_text: str | None) -> dict[str, Any]:
|
||||
"""
|
||||
解析错误响应为结构化数据
|
||||
|
||||
@@ -265,7 +265,7 @@ class ErrorClassifier:
|
||||
|
||||
return result
|
||||
|
||||
def is_client_error(self, error_text: Optional[str]) -> bool:
|
||||
def is_client_error(self, error_text: str | None) -> bool:
|
||||
"""
|
||||
检测错误响应是否为客户端错误(不应重试)
|
||||
|
||||
@@ -301,7 +301,7 @@ class ErrorClassifier:
|
||||
search_text = f"{parsed['message']} {parsed['raw']}".lower()
|
||||
return any(pattern.lower() in search_text for pattern in self.CLIENT_ERROR_PATTERNS)
|
||||
|
||||
def _is_compatibility_error(self, error_text: Optional[str]) -> bool:
|
||||
def _is_compatibility_error(self, error_text: str | None) -> bool:
|
||||
"""
|
||||
检测错误响应是否为 Provider 兼容性错误(应触发故障转移)
|
||||
|
||||
@@ -320,7 +320,7 @@ class ErrorClassifier:
|
||||
search_text = error_text.lower()
|
||||
return any(pattern.lower() in search_text for pattern in self.COMPATIBILITY_ERROR_PATTERNS)
|
||||
|
||||
def _is_thinking_error(self, error_text: Optional[str]) -> bool:
|
||||
def _is_thinking_error(self, error_text: str | None) -> bool:
|
||||
"""
|
||||
检测错误响应是否为 Thinking 块相关错误(签名错误或结构错误)
|
||||
|
||||
@@ -339,7 +339,7 @@ class ErrorClassifier:
|
||||
search_text = error_text.lower()
|
||||
return any(p.lower() in search_text for p in self.THINKING_ERROR_PATTERNS)
|
||||
|
||||
def _extract_error_message(self, error_text: Optional[str]) -> Optional[str]:
|
||||
def _extract_error_message(self, error_text: str | None) -> str | None:
|
||||
"""
|
||||
从错误响应中提取错误消息
|
||||
|
||||
@@ -404,9 +404,9 @@ class ErrorClassifier:
|
||||
self,
|
||||
key: ProviderAPIKey,
|
||||
provider_name: str,
|
||||
current_rpm: Optional[int],
|
||||
current_rpm: int | None,
|
||||
exception: ProviderRateLimitException,
|
||||
request_id: Optional[str] = None,
|
||||
request_id: str | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
处理 429 速率限制错误的自适应调整
|
||||
@@ -468,8 +468,8 @@ class ErrorClassifier:
|
||||
self,
|
||||
error: httpx.HTTPStatusError,
|
||||
provider_name: str,
|
||||
error_response_text: Optional[str] = None,
|
||||
) -> Union[ProviderException, UpstreamClientException]:
|
||||
error_response_text: str | None = None,
|
||||
) -> ProviderException | UpstreamClientException:
|
||||
"""
|
||||
转换 HTTP 错误为 Provider 异常
|
||||
|
||||
@@ -559,14 +559,14 @@ class ErrorClassifier:
|
||||
endpoint: ProviderEndpoint,
|
||||
key: ProviderAPIKey,
|
||||
affinity_key: str,
|
||||
api_format: Union[str, APIFormat],
|
||||
api_format: str | APIFormat,
|
||||
global_model_id: str,
|
||||
request_id: Optional[str],
|
||||
captured_key_concurrent: Optional[int],
|
||||
elapsed_ms: Optional[int],
|
||||
request_id: str | None,
|
||||
captured_key_concurrent: int | None,
|
||||
elapsed_ms: int | None,
|
||||
attempt: int,
|
||||
max_attempts: int,
|
||||
) -> Dict[str, Any]:
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
处理 HTTP 错误,返回 extra_data
|
||||
|
||||
@@ -609,7 +609,7 @@ class ErrorClassifier:
|
||||
converted_error = self.convert_http_error(http_error, provider_name, error_response_text)
|
||||
|
||||
# 构建 extra_data,包含转换后的异常
|
||||
extra_data: Dict[str, Any] = {
|
||||
extra_data: dict[str, Any] = {
|
||||
"converted_error": converted_error,
|
||||
}
|
||||
if error_response_text:
|
||||
@@ -702,11 +702,11 @@ class ErrorClassifier:
|
||||
endpoint: ProviderEndpoint,
|
||||
key: ProviderAPIKey,
|
||||
affinity_key: str,
|
||||
api_format: Union[str, APIFormat],
|
||||
api_format: str | APIFormat,
|
||||
global_model_id: str,
|
||||
captured_key_concurrent: Optional[int],
|
||||
elapsed_ms: Optional[int],
|
||||
request_id: Optional[str],
|
||||
captured_key_concurrent: int | None,
|
||||
elapsed_ms: int | None,
|
||||
request_id: str | None,
|
||||
attempt: int,
|
||||
max_attempts: int,
|
||||
) -> None:
|
||||
|
||||
@@ -21,9 +21,10 @@
|
||||
- 本类作为协调者,组合使用上述组件
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Callable, Dict, List, NoReturn, Optional, Tuple, Union
|
||||
from typing import Any, NoReturn
|
||||
|
||||
from collections.abc import Callable
|
||||
|
||||
import httpx
|
||||
from redis import Redis
|
||||
@@ -80,7 +81,7 @@ class FallbackOrchestrator:
|
||||
- 优势:可预测、高效、公平、资源友好
|
||||
"""
|
||||
|
||||
def __init__(self, db: Session, redis_client: Optional[Redis] = None) -> None:
|
||||
def __init__(self, db: Session, redis_client: Redis | None = None) -> None:
|
||||
"""
|
||||
初始化编排器
|
||||
|
||||
@@ -90,15 +91,15 @@ class FallbackOrchestrator:
|
||||
"""
|
||||
self.db = db
|
||||
self.redis = redis_client
|
||||
self.cache_scheduler: Optional[CacheAwareScheduler] = None
|
||||
self.cache_scheduler: CacheAwareScheduler | None = None
|
||||
self.concurrency_manager: Any = None
|
||||
self.adaptive_manager = get_adaptive_rpm_manager() # 自适应 RPM 管理器
|
||||
self.request_executor: Optional[RequestExecutor] = None
|
||||
self.request_executor: RequestExecutor | None = None
|
||||
|
||||
# 拆分后的组件(延迟初始化)
|
||||
self._candidate_resolver: Optional[CandidateResolver] = None
|
||||
self._request_dispatcher: Optional[RequestDispatcher] = None
|
||||
self._error_classifier: Optional[ErrorClassifier] = None
|
||||
self._candidate_resolver: CandidateResolver | None = None
|
||||
self._request_dispatcher: RequestDispatcher | None = None
|
||||
self._error_classifier: ErrorClassifier | None = None
|
||||
|
||||
async def _ensure_initialized(self) -> None:
|
||||
"""确保异步组件已初始化"""
|
||||
@@ -172,11 +173,11 @@ class FallbackOrchestrator:
|
||||
api_format: APIFormat,
|
||||
model_name: str,
|
||||
affinity_key: str,
|
||||
user_api_key: Optional[ApiKey] = None,
|
||||
request_id: Optional[str] = None,
|
||||
user_api_key: ApiKey | None = None,
|
||||
request_id: str | None = None,
|
||||
is_stream: bool = False,
|
||||
capability_requirements: Optional[Dict[str, bool]] = None,
|
||||
) -> Tuple[List[ProviderCandidate], str]:
|
||||
capability_requirements: dict[str, bool] | None = None,
|
||||
) -> tuple[list[ProviderCandidate], str]:
|
||||
"""
|
||||
收集所有可用的 Provider/Endpoint/Key 候选组合
|
||||
|
||||
@@ -210,12 +211,12 @@ class FallbackOrchestrator:
|
||||
|
||||
def _create_candidate_records(
|
||||
self,
|
||||
all_candidates: List[ProviderCandidate],
|
||||
request_id: Optional[str],
|
||||
all_candidates: list[ProviderCandidate],
|
||||
request_id: str | None,
|
||||
user_id: str,
|
||||
user_api_key: ApiKey,
|
||||
required_capabilities: Optional[Dict[str, bool]] = None,
|
||||
) -> Dict[Tuple[int, int], str]:
|
||||
required_capabilities: dict[str, bool] | None = None,
|
||||
) -> dict[tuple[int, int], str]:
|
||||
"""
|
||||
为所有候选预先创建 available 状态记录(批量插入优化)
|
||||
|
||||
@@ -248,7 +249,7 @@ class FallbackOrchestrator:
|
||||
candidate_record_id: str,
|
||||
user_api_key: ApiKey,
|
||||
request_func: Callable[..., Any],
|
||||
request_id: Optional[str],
|
||||
request_id: str | None,
|
||||
api_format: APIFormat,
|
||||
model_name: str,
|
||||
affinity_key: str,
|
||||
@@ -256,7 +257,7 @@ class FallbackOrchestrator:
|
||||
attempt_counter: int,
|
||||
max_attempts: int,
|
||||
is_stream: bool = False,
|
||||
) -> Tuple[Any, str, str, str, str, str]:
|
||||
) -> tuple[Any, str, str, str, str, str]:
|
||||
"""
|
||||
尝试单个候选执行请求
|
||||
|
||||
@@ -305,12 +306,12 @@ class FallbackOrchestrator:
|
||||
def _handle_thinking_signature_error(
|
||||
self,
|
||||
converted_error: ThinkingSignatureException,
|
||||
request_id: Optional[str],
|
||||
request_id: str | None,
|
||||
candidate_record_id: str,
|
||||
elapsed_ms: int,
|
||||
captured_key_concurrent: Optional[int],
|
||||
serializable_extra_data: Dict[str, Any],
|
||||
request_body_ref: Optional[Dict[str, Any]],
|
||||
captured_key_concurrent: int | None,
|
||||
serializable_extra_data: dict[str, Any],
|
||||
request_body_ref: dict[str, Any] | None,
|
||||
) -> str:
|
||||
"""
|
||||
处理 ThinkingSignatureException 错误
|
||||
@@ -395,8 +396,8 @@ class FallbackOrchestrator:
|
||||
candidate_record_id: str,
|
||||
error: ThinkingSignatureException,
|
||||
elapsed_ms: int,
|
||||
captured_key_concurrent: Optional[int],
|
||||
extra_data: Dict[str, Any],
|
||||
captured_key_concurrent: int | None,
|
||||
extra_data: dict[str, Any],
|
||||
) -> None:
|
||||
"""标记 Thinking 签名错误导致的候选失败"""
|
||||
RequestCandidateService.mark_candidate_failed(
|
||||
@@ -420,10 +421,10 @@ class FallbackOrchestrator:
|
||||
affinity_key: str,
|
||||
api_format: APIFormat,
|
||||
global_model_id: str,
|
||||
request_id: Optional[str],
|
||||
request_id: str | None,
|
||||
attempt: int,
|
||||
max_attempts: int,
|
||||
request_body_ref: Optional[Dict[str, Any]] = None,
|
||||
request_body_ref: dict[str, Any] | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
处理候选执行错误
|
||||
@@ -651,7 +652,7 @@ class FallbackOrchestrator:
|
||||
|
||||
def _create_pending_usage_record(
|
||||
self,
|
||||
request_id: Optional[str],
|
||||
request_id: str | None,
|
||||
user_api_key: ApiKey,
|
||||
model_name: str,
|
||||
is_stream: bool,
|
||||
@@ -682,23 +683,23 @@ class FallbackOrchestrator:
|
||||
|
||||
async def _execute_candidates_loop(
|
||||
self,
|
||||
all_candidates: List[ProviderCandidate],
|
||||
candidate_record_map: Dict[Tuple[int, int], str],
|
||||
all_candidates: list[ProviderCandidate],
|
||||
candidate_record_map: dict[tuple[int, int], str],
|
||||
user_api_key: ApiKey,
|
||||
request_func: Callable[..., Any],
|
||||
request_id: Optional[str],
|
||||
request_id: str | None,
|
||||
api_format_enum: APIFormat,
|
||||
model_name: str,
|
||||
affinity_key: str,
|
||||
global_model_id: str,
|
||||
is_stream: bool = False,
|
||||
request_body_ref: Optional[Dict[str, Any]] = None,
|
||||
) -> Tuple[Any, str, Optional[str], Optional[str], Optional[str], Optional[str]]:
|
||||
request_body_ref: dict[str, Any] | None = None,
|
||||
) -> tuple[Any, str, str | None, str | None, str | None, str | None]:
|
||||
"""遍历所有候选执行请求,返回第一个成功的结果或抛出异常"""
|
||||
attempt_counter = 0
|
||||
max_attempts = 0
|
||||
last_error: Optional[Exception] = None
|
||||
last_candidate: Optional[ProviderCandidate] = None
|
||||
last_error: Exception | None = None
|
||||
last_candidate: ProviderCandidate | None = None
|
||||
|
||||
for candidate_index, candidate in enumerate(all_candidates):
|
||||
last_candidate = candidate
|
||||
@@ -728,8 +729,8 @@ class FallbackOrchestrator:
|
||||
)
|
||||
|
||||
if result["success"]:
|
||||
response: Tuple[
|
||||
Any, str, Optional[str], Optional[str], Optional[str], Optional[str]
|
||||
response: tuple[
|
||||
Any, str, str | None, str | None, str | None, str | None
|
||||
] = result["response"]
|
||||
return response
|
||||
|
||||
@@ -753,10 +754,10 @@ class FallbackOrchestrator:
|
||||
self,
|
||||
candidate: ProviderCandidate,
|
||||
candidate_index: int,
|
||||
candidate_record_map: Dict[Tuple[int, int], str],
|
||||
candidate_record_map: dict[tuple[int, int], str],
|
||||
user_api_key: ApiKey,
|
||||
request_func: Callable[..., Any],
|
||||
request_id: Optional[str],
|
||||
request_id: str | None,
|
||||
api_format_enum: APIFormat,
|
||||
model_name: str,
|
||||
affinity_key: str,
|
||||
@@ -764,14 +765,14 @@ class FallbackOrchestrator:
|
||||
attempt_counter: int,
|
||||
max_attempts: int,
|
||||
is_stream: bool = False,
|
||||
request_body_ref: Optional[Dict[str, Any]] = None,
|
||||
) -> Dict[str, Any]:
|
||||
request_body_ref: dict[str, Any] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""尝试单个候选(含重试逻辑),返回执行结果"""
|
||||
provider = candidate.provider
|
||||
endpoint = candidate.endpoint
|
||||
# 从 Provider 读取 max_retries(已从 Endpoint 迁移)
|
||||
max_retries_for_candidate = int(provider.max_retries or 2) if candidate.is_cached else 1
|
||||
last_error: Optional[Exception] = None
|
||||
last_error: Exception | None = None
|
||||
|
||||
retry_index = 0
|
||||
while retry_index < max_retries_for_candidate:
|
||||
@@ -876,8 +877,8 @@ class FallbackOrchestrator:
|
||||
|
||||
def _attach_metadata_to_error(
|
||||
self,
|
||||
error: Optional[Exception],
|
||||
candidate: Optional[ProviderCandidate],
|
||||
error: Exception | None,
|
||||
candidate: ProviderCandidate | None,
|
||||
model_name: str,
|
||||
api_format_enum: APIFormat,
|
||||
) -> None:
|
||||
@@ -915,12 +916,12 @@ class FallbackOrchestrator:
|
||||
|
||||
def _raise_all_failed_exception(
|
||||
self,
|
||||
request_id: Optional[str],
|
||||
request_id: str | None,
|
||||
max_attempts: int,
|
||||
last_candidate: Optional[ProviderCandidate],
|
||||
last_candidate: ProviderCandidate | None,
|
||||
model_name: str,
|
||||
api_format_enum: APIFormat,
|
||||
last_error: Optional[Exception] = None,
|
||||
last_error: Exception | None = None,
|
||||
) -> NoReturn:
|
||||
"""所有组合都失败时抛出异常"""
|
||||
logger.error(f" [{request_id}] 所有 {max_attempts} 个组合均失败")
|
||||
@@ -937,8 +938,8 @@ class FallbackOrchestrator:
|
||||
}
|
||||
|
||||
# 提取上游错误响应
|
||||
upstream_status: Optional[int] = None
|
||||
upstream_response: Optional[str] = None
|
||||
upstream_status: int | None = None
|
||||
upstream_response: str | None = None
|
||||
if last_error:
|
||||
# 从 httpx.HTTPStatusError 提取
|
||||
if isinstance(last_error, httpx.HTTPStatusError):
|
||||
@@ -981,15 +982,15 @@ class FallbackOrchestrator:
|
||||
|
||||
async def execute_with_fallback(
|
||||
self,
|
||||
api_format: Union[str, APIFormat],
|
||||
api_format: str | APIFormat,
|
||||
model_name: str,
|
||||
user_api_key: ApiKey,
|
||||
request_func: Callable[[Provider, ProviderEndpoint, ProviderAPIKey], Any],
|
||||
request_id: Optional[str] = None,
|
||||
request_id: str | None = None,
|
||||
is_stream: bool = False,
|
||||
capability_requirements: Optional[Dict[str, bool]] = None,
|
||||
request_body_ref: Optional[Dict[str, Any]] = None,
|
||||
) -> Tuple[Any, str, Optional[str], Optional[str], Optional[str], Optional[str]]:
|
||||
capability_requirements: dict[str, bool] | None = None,
|
||||
request_body_ref: dict[str, Any] | None = None,
|
||||
) -> tuple[Any, str, str | None, str | None, str | None, str | None]:
|
||||
"""
|
||||
执行请求,并在失败时自动故障转移(缓存感知)
|
||||
|
||||
|
||||
@@ -4,7 +4,9 @@
|
||||
负责执行单个候选请求
|
||||
"""
|
||||
|
||||
from typing import Any, Callable, Optional, Tuple
|
||||
from typing import Any
|
||||
|
||||
from collections.abc import Callable
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
@@ -31,7 +33,7 @@ class RequestDispatcher:
|
||||
self,
|
||||
db: Session,
|
||||
request_executor: RequestExecutor,
|
||||
cache_scheduler: Optional[CacheAwareScheduler] = None,
|
||||
cache_scheduler: CacheAwareScheduler | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
初始化请求分发器
|
||||
@@ -53,7 +55,7 @@ class RequestDispatcher:
|
||||
candidate_record_id: str,
|
||||
user_api_key: ApiKey,
|
||||
request_func: Callable[..., Any],
|
||||
request_id: Optional[str],
|
||||
request_id: str | None,
|
||||
api_format: APIFormat,
|
||||
model_name: str,
|
||||
affinity_key: str,
|
||||
@@ -61,7 +63,7 @@ class RequestDispatcher:
|
||||
attempt_counter: int,
|
||||
max_attempts: int,
|
||||
is_stream: bool = False,
|
||||
) -> Tuple[Any, str, str, str, str, str]:
|
||||
) -> tuple[Any, str, str, str, str, str]:
|
||||
"""
|
||||
执行请求并返回结果
|
||||
|
||||
@@ -98,7 +100,7 @@ class RequestDispatcher:
|
||||
key_id = str(key.id)
|
||||
cache_ttl_minutes = int(key.cache_ttl_minutes or 0)
|
||||
provider_supports_caching = cache_ttl_minutes > 0
|
||||
provider_cache_ttl_seconds: Optional[int] = (
|
||||
provider_cache_ttl_seconds: int | None = (
|
||||
cache_ttl_minutes * 60 if cache_ttl_minutes > 0 else None
|
||||
)
|
||||
|
||||
|
||||
@@ -2,15 +2,13 @@
|
||||
API 格式辅助函数,确保在调度/编排链路中使用统一的枚举值。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Optional, Union
|
||||
|
||||
from src.core.api_format import APIFormat, resolve_api_format
|
||||
|
||||
|
||||
def normalize_api_format(
|
||||
value: Union[str, APIFormat, None], default: APIFormat = APIFormat.CLAUDE
|
||||
value: str | APIFormat | None, default: APIFormat = APIFormat.CLAUDE
|
||||
) -> APIFormat:
|
||||
"""
|
||||
将任意字符串/枚举值归一化为 APIFormat。
|
||||
|
||||
@@ -3,7 +3,6 @@
|
||||
负责提供商选择、模型映射和请求处理
|
||||
"""
|
||||
|
||||
from typing import Dict
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
@@ -92,7 +91,7 @@ class ProviderService:
|
||||
|
||||
def calculate_cost(
|
||||
self, provider: Provider, model: str, input_tokens: int, output_tokens: int
|
||||
) -> Dict[str, float]:
|
||||
) -> dict[str, float]:
|
||||
"""
|
||||
计算使用成本
|
||||
|
||||
@@ -107,7 +106,7 @@ class ProviderService:
|
||||
"""
|
||||
return self.mapper.calculate_cost(model, provider.id, input_tokens, output_tokens)
|
||||
|
||||
def get_available_models(self) -> Dict[str, list]:
|
||||
def get_available_models(self) -> dict[str, list]:
|
||||
"""
|
||||
获取所有可用的模型
|
||||
|
||||
|
||||
@@ -7,7 +7,7 @@
|
||||
"""
|
||||
|
||||
import re
|
||||
from typing import TYPE_CHECKING, Any, Dict, Optional
|
||||
from typing import TYPE_CHECKING, Any
|
||||
from urllib.parse import urlencode
|
||||
|
||||
from src.core.api_format import APIFormat, get_default_path, resolve_api_format
|
||||
@@ -64,10 +64,10 @@ def _normalize_base_url(base_url: str, path: str) -> str:
|
||||
|
||||
|
||||
def build_provider_url(
|
||||
endpoint: "ProviderEndpoint",
|
||||
endpoint: ProviderEndpoint,
|
||||
*,
|
||||
query_params: Optional[Dict[str, Any]] = None,
|
||||
path_params: Optional[Dict[str, Any]] = None,
|
||||
query_params: dict[str, Any] | None = None,
|
||||
path_params: dict[str, Any] | None = None,
|
||||
is_stream: bool = False,
|
||||
) -> str:
|
||||
"""
|
||||
@@ -142,7 +142,7 @@ def build_provider_url(
|
||||
return url
|
||||
|
||||
|
||||
def _resolve_default_path(api_format: Optional[str]) -> str:
|
||||
def _resolve_default_path(api_format: str | None) -> str:
|
||||
"""
|
||||
根据 API 格式返回默认路径
|
||||
"""
|
||||
|
||||
@@ -3,7 +3,7 @@ Anyrouter 余额查询操作(含自动签到)
|
||||
"""
|
||||
|
||||
import time
|
||||
from typing import Any, Dict, Optional
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
@@ -101,7 +101,7 @@ class AnyrouterBalanceAction(BalanceAction):
|
||||
)
|
||||
|
||||
def _handle_http_error(
|
||||
self, response: httpx.Response, raw_data: Optional[Dict[str, Any]] = None
|
||||
self, response: httpx.Response, raw_data: dict[str, Any] | None = None
|
||||
) -> ActionResult:
|
||||
"""处理 HTTP 错误响应"""
|
||||
status_code = response.status_code
|
||||
@@ -117,7 +117,7 @@ class AnyrouterBalanceAction(BalanceAction):
|
||||
|
||||
return super()._handle_http_error(response, raw_data)
|
||||
|
||||
async def _do_checkin(self, client: httpx.AsyncClient) -> Optional[Dict[str, Any]]:
|
||||
async def _do_checkin(self, client: httpx.AsyncClient) -> dict[str, Any] | None:
|
||||
"""
|
||||
执行自动签到(始终执行)
|
||||
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
"""
|
||||
|
||||
from abc import abstractmethod
|
||||
from typing import Any, Dict, Optional
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
@@ -80,7 +80,7 @@ class BalanceAction(ProviderAction):
|
||||
"""
|
||||
pass
|
||||
|
||||
async def _do_checkin(self, client: httpx.AsyncClient) -> Optional[Dict[str, Any]]:
|
||||
async def _do_checkin(self, client: httpx.AsyncClient) -> dict[str, Any] | None:
|
||||
"""
|
||||
执行签到(子类可选实现)
|
||||
|
||||
@@ -97,11 +97,11 @@ class BalanceAction(ProviderAction):
|
||||
|
||||
def _create_balance_info(
|
||||
self,
|
||||
total_granted: Optional[float] = None,
|
||||
total_used: Optional[float] = None,
|
||||
total_available: Optional[float] = None,
|
||||
total_granted: float | None = None,
|
||||
total_used: float | None = None,
|
||||
total_available: float | None = None,
|
||||
currency: str = "USD",
|
||||
extra: Optional[Dict[str, Any]] = None,
|
||||
extra: dict[str, Any] | None = None,
|
||||
) -> BalanceInfo:
|
||||
"""
|
||||
创建余额信息对象
|
||||
@@ -135,7 +135,7 @@ class BalanceAction(ProviderAction):
|
||||
extra=extra if extra is not None else {},
|
||||
)
|
||||
|
||||
def _to_float(self, value: Any) -> Optional[float]:
|
||||
def _to_float(self, value: Any) -> float | None:
|
||||
"""转换为浮点数"""
|
||||
if value is None:
|
||||
return None
|
||||
@@ -145,7 +145,7 @@ class BalanceAction(ProviderAction):
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def get_config_schema(cls) -> Dict[str, Any]:
|
||||
def get_config_schema(cls) -> dict[str, Any]:
|
||||
"""获取操作配置 schema(子类可重写)"""
|
||||
return {
|
||||
"type": "object",
|
||||
|
||||
@@ -3,8 +3,7 @@ Provider 操作抽象基类
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Dict, Optional
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
@@ -30,7 +29,7 @@ class ProviderAction(ABC):
|
||||
# 默认缓存时间(秒)
|
||||
default_cache_ttl: int = 300
|
||||
|
||||
def __init__(self, config: Optional[Dict[str, Any]] = None):
|
||||
def __init__(self, config: dict[str, Any] | None = None):
|
||||
"""
|
||||
初始化操作
|
||||
|
||||
@@ -52,7 +51,7 @@ class ProviderAction(ABC):
|
||||
"""
|
||||
pass
|
||||
|
||||
def _extract_field(self, data: Any, path: Optional[str]) -> Any:
|
||||
def _extract_field(self, data: Any, path: str | None) -> Any:
|
||||
"""
|
||||
从响应数据中提取字段
|
||||
|
||||
@@ -86,9 +85,9 @@ class ProviderAction(ABC):
|
||||
def _make_success_result(
|
||||
self,
|
||||
data: Any = None,
|
||||
message: Optional[str] = None,
|
||||
response_time_ms: Optional[int] = None,
|
||||
raw_response: Optional[Dict[str, Any]] = None,
|
||||
message: str | None = None,
|
||||
response_time_ms: int | None = None,
|
||||
raw_response: dict[str, Any] | None = None,
|
||||
) -> ActionResult:
|
||||
"""创建成功结果"""
|
||||
return ActionResult(
|
||||
@@ -104,9 +103,9 @@ class ProviderAction(ABC):
|
||||
def _make_error_result(
|
||||
self,
|
||||
status: ActionStatus,
|
||||
message: Optional[str] = None,
|
||||
retry_after_seconds: Optional[int] = None,
|
||||
raw_response: Optional[Dict[str, Any]] = None,
|
||||
message: str | None = None,
|
||||
retry_after_seconds: int | None = None,
|
||||
raw_response: dict[str, Any] | None = None,
|
||||
) -> ActionResult:
|
||||
"""创建错误结果"""
|
||||
return ActionResult(
|
||||
@@ -119,7 +118,7 @@ class ProviderAction(ABC):
|
||||
)
|
||||
|
||||
def _handle_http_error(
|
||||
self, response: httpx.Response, raw_data: Optional[Dict[str, Any]] = None
|
||||
self, response: httpx.Response, raw_data: dict[str, Any] | None = None
|
||||
) -> ActionResult:
|
||||
"""处理 HTTP 错误响应"""
|
||||
status_code = response.status_code
|
||||
@@ -153,7 +152,7 @@ class ProviderAction(ABC):
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_config_schema(cls) -> Dict[str, Any]:
|
||||
def get_config_schema(cls) -> dict[str, Any]:
|
||||
"""
|
||||
获取操作配置 JSON Schema(用于前端表单生成)
|
||||
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
"""
|
||||
|
||||
from abc import abstractmethod
|
||||
from typing import Any, Dict, Optional
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
@@ -44,10 +44,10 @@ class CheckinAction(ProviderAction):
|
||||
|
||||
def _create_checkin_info(
|
||||
self,
|
||||
reward: Optional[float] = None,
|
||||
streak_days: Optional[int] = None,
|
||||
message: Optional[str] = None,
|
||||
extra: Optional[Dict[str, Any]] = None,
|
||||
reward: float | None = None,
|
||||
streak_days: int | None = None,
|
||||
message: str | None = None,
|
||||
extra: dict[str, Any] | None = None,
|
||||
) -> CheckinInfo:
|
||||
"""
|
||||
创建签到信息对象
|
||||
@@ -71,7 +71,7 @@ class CheckinAction(ProviderAction):
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_config_schema(cls) -> Dict[str, Any]:
|
||||
def get_config_schema(cls) -> dict[str, Any]:
|
||||
"""获取操作配置 schema(子类可重写)"""
|
||||
return {
|
||||
"type": "object",
|
||||
|
||||
@@ -3,7 +3,7 @@ Cubence 余额查询操作
|
||||
"""
|
||||
|
||||
import time
|
||||
from typing import Any, Dict, Optional
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
@@ -85,7 +85,7 @@ class CubenceBalanceAction(BalanceAction):
|
||||
)
|
||||
|
||||
def _handle_http_error(
|
||||
self, response: httpx.Response, raw_data: Optional[Dict[str, Any]] = None
|
||||
self, response: httpx.Response, raw_data: dict[str, Any] | None = None
|
||||
) -> ActionResult:
|
||||
"""处理 HTTP 错误响应(Cubence 专用)"""
|
||||
status_code = response.status_code
|
||||
@@ -117,7 +117,7 @@ class CubenceBalanceAction(BalanceAction):
|
||||
charity_balance = balance_data.get("charity_balance_dollar")
|
||||
|
||||
# 窗口限额信息
|
||||
extra: Dict[str, Any] = {}
|
||||
extra: dict[str, Any] = {}
|
||||
|
||||
# 5小时窗口限额
|
||||
five_hour = subscription_limits.get("five_hour", {})
|
||||
|
||||
@@ -4,7 +4,7 @@ NekoCode 余额查询操作
|
||||
|
||||
import time
|
||||
from datetime import datetime
|
||||
from typing import Any, Dict, Optional
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
@@ -122,7 +122,7 @@ class NekoCodeBalanceAction(BalanceAction):
|
||||
logger.debug(f"解析 effective_start_date 失败: {e}")
|
||||
|
||||
# 构建 extra 信息
|
||||
extra: Dict[str, Any] = {
|
||||
extra: dict[str, Any] = {
|
||||
"plan_name": plan_name,
|
||||
"subscription_status": status,
|
||||
"daily_quota_limit": daily_quota_limit,
|
||||
@@ -161,7 +161,7 @@ class NekoCodeBalanceAction(BalanceAction):
|
||||
)
|
||||
|
||||
def _handle_http_error(
|
||||
self, response: httpx.Response, raw_data: Optional[Dict[str, Any]] = None
|
||||
self, response: httpx.Response, raw_data: dict[str, Any] | None = None
|
||||
) -> ActionResult:
|
||||
"""处理 HTTP 错误响应"""
|
||||
status_code = response.status_code
|
||||
@@ -178,7 +178,7 @@ class NekoCodeBalanceAction(BalanceAction):
|
||||
return super()._handle_http_error(response, raw_data)
|
||||
|
||||
@classmethod
|
||||
def get_config_schema(cls) -> Dict[str, Any]:
|
||||
def get_config_schema(cls) -> dict[str, Any]:
|
||||
"""获取操作配置 schema"""
|
||||
return {
|
||||
"type": "object",
|
||||
|
||||
@@ -3,7 +3,7 @@ New API 余额查询操作
|
||||
"""
|
||||
|
||||
import time
|
||||
from typing import Any, Dict, Optional
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
@@ -111,7 +111,7 @@ class NewApiBalanceAction(BalanceAction):
|
||||
currency=self.config.get("currency", "USD"),
|
||||
)
|
||||
|
||||
async def _do_checkin(self, client: httpx.AsyncClient) -> Optional[Dict[str, Any]]:
|
||||
async def _do_checkin(self, client: httpx.AsyncClient) -> dict[str, Any] | None:
|
||||
"""
|
||||
执行签到(静默,不抛出异常)
|
||||
|
||||
@@ -184,7 +184,7 @@ class NewApiBalanceAction(BalanceAction):
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def get_config_schema(cls) -> Dict[str, Any]:
|
||||
def get_config_schema(cls) -> dict[str, Any]:
|
||||
"""获取操作配置 schema"""
|
||||
return {
|
||||
"type": "object",
|
||||
|
||||
@@ -4,7 +4,7 @@ YesCode 余额查询操作
|
||||
|
||||
import asyncio
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Any, Dict
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
@@ -15,7 +15,7 @@ from src.services.provider_ops.types import ActionResult, ActionStatus, BalanceI
|
||||
async def fetch_yescode_combined_data(
|
||||
client: httpx.AsyncClient,
|
||||
base_url: str,
|
||||
) -> Dict[str, Any]:
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
获取 YesCode 合并数据(balance + profile)
|
||||
|
||||
@@ -31,7 +31,7 @@ async def fetch_yescode_combined_data(
|
||||
合并后的数据字典
|
||||
"""
|
||||
base_url = base_url.rstrip("/")
|
||||
result: Dict[str, Any] = {}
|
||||
result: dict[str, Any] = {}
|
||||
|
||||
# 并发调用两个接口
|
||||
balance_task = client.get(f"{base_url}/api/v1/user/balance")
|
||||
@@ -42,7 +42,7 @@ async def fetch_yescode_combined_data(
|
||||
)
|
||||
|
||||
# 解析 balance 接口
|
||||
balance_data: Dict[str, Any] = {}
|
||||
balance_data: dict[str, Any] = {}
|
||||
if isinstance(balance_resp, httpx.Response) and balance_resp.status_code == 200:
|
||||
try:
|
||||
balance_data = balance_resp.json()
|
||||
@@ -51,7 +51,7 @@ async def fetch_yescode_combined_data(
|
||||
pass
|
||||
|
||||
# 解析 profile 接口
|
||||
profile_data: Dict[str, Any] = {}
|
||||
profile_data: dict[str, Any] = {}
|
||||
if isinstance(profile_resp, httpx.Response) and profile_resp.status_code == 200:
|
||||
try:
|
||||
profile_data = profile_resp.json()
|
||||
@@ -87,7 +87,7 @@ async def fetch_yescode_combined_data(
|
||||
return result
|
||||
|
||||
|
||||
def parse_yescode_balance_extra(data: Dict[str, Any]) -> Dict[str, Any]:
|
||||
def parse_yescode_balance_extra(data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
解析 YesCode 余额额外信息
|
||||
|
||||
@@ -97,7 +97,7 @@ def parse_yescode_balance_extra(data: Dict[str, Any]) -> Dict[str, Any]:
|
||||
Returns:
|
||||
统一格式的 extra 字典
|
||||
"""
|
||||
extra: Dict[str, Any] = {}
|
||||
extra: dict[str, Any] = {}
|
||||
|
||||
pay_as_you_go = data.get("pay_as_you_go_balance", 0)
|
||||
subscription = data.get("subscription_balance", 0)
|
||||
|
||||
@@ -6,7 +6,7 @@ Anyrouter 架构
|
||||
|
||||
import base64
|
||||
import re
|
||||
from typing import Any, Dict, List, Optional, Tuple, Type
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
@@ -118,7 +118,7 @@ def _extract_session_from_cookie(cookie_string: str) -> str:
|
||||
return cookie_string.strip()
|
||||
|
||||
|
||||
def _parse_session_user_id(cookie_input: str) -> Tuple[Optional[str], Optional[str]]:
|
||||
def _parse_session_user_id(cookie_input: str) -> tuple[str | None, str | None]:
|
||||
"""
|
||||
从 session cookie 中解析用户 ID 和用户名
|
||||
|
||||
@@ -213,8 +213,8 @@ def _parse_session_user_id(cookie_input: str) -> Tuple[Optional[str], Optional[s
|
||||
|
||||
|
||||
async def _get_acw_cookie(
|
||||
base_url: str, timeout: float = 10, proxy: Optional[str] = None
|
||||
) -> Optional[str]:
|
||||
base_url: str, timeout: float = 10, proxy: str | None = None
|
||||
) -> str | None:
|
||||
"""
|
||||
获取 acw_sc__v2 Cookie
|
||||
|
||||
@@ -230,7 +230,7 @@ async def _get_acw_cookie(
|
||||
"""
|
||||
try:
|
||||
# 构建 client 参数
|
||||
client_kwargs: Dict[str, Any] = {
|
||||
client_kwargs: dict[str, Any] = {
|
||||
"timeout": timeout,
|
||||
"verify": get_ssl_context(),
|
||||
}
|
||||
@@ -278,13 +278,13 @@ class AnyrouterConnector(ProviderConnector):
|
||||
auth_type = ConnectorAuthType.COOKIE
|
||||
display_name = "Anyrouter Cookie"
|
||||
|
||||
def __init__(self, base_url: str, config: Optional[Dict[str, Any]] = None):
|
||||
def __init__(self, base_url: str, config: dict[str, Any] | None = None):
|
||||
super().__init__(base_url, config)
|
||||
self._session_cookie: Optional[str] = None
|
||||
self._acw_cookie: Optional[str] = None
|
||||
self._user_id: Optional[str] = None
|
||||
self._session_cookie: str | None = None
|
||||
self._acw_cookie: str | None = None
|
||||
self._user_id: str | None = None
|
||||
|
||||
async def connect(self, credentials: Dict[str, Any]) -> bool:
|
||||
async def connect(self, credentials: dict[str, Any]) -> bool:
|
||||
"""建立连接"""
|
||||
session_cookie = credentials.get("session_cookie")
|
||||
if not session_cookie:
|
||||
@@ -336,7 +336,7 @@ class AnyrouterConnector(ProviderConnector):
|
||||
return request
|
||||
|
||||
@classmethod
|
||||
def get_credentials_schema(cls) -> Dict[str, Any]:
|
||||
def get_credentials_schema(cls) -> dict[str, Any]:
|
||||
"""获取凭据配置 schema"""
|
||||
return {
|
||||
"type": "object",
|
||||
@@ -368,13 +368,13 @@ class AnyrouterArchitecture(ProviderArchitecture):
|
||||
display_name = "Anyrouter"
|
||||
description = "Anyrouter 中转站预设配置,使用 Cookie 认证"
|
||||
|
||||
supported_connectors: List[Type[ProviderConnector]] = [
|
||||
supported_connectors: list[type[ProviderConnector]] = [
|
||||
AnyrouterConnector,
|
||||
]
|
||||
|
||||
supported_actions: List[Type[ProviderAction]] = [AnyrouterBalanceAction]
|
||||
supported_actions: list[type[ProviderAction]] = [AnyrouterBalanceAction]
|
||||
|
||||
default_action_configs: Dict[ProviderActionType, Dict[str, Any]] = {
|
||||
default_action_configs: dict[ProviderActionType, dict[str, Any]] = {
|
||||
ProviderActionType.QUERY_BALANCE: {
|
||||
"endpoint": "/api/user/self",
|
||||
"method": "GET",
|
||||
@@ -383,7 +383,7 @@ class AnyrouterArchitecture(ProviderArchitecture):
|
||||
},
|
||||
}
|
||||
|
||||
def get_credentials_schema(self) -> Dict[str, Any]:
|
||||
def get_credentials_schema(self) -> dict[str, Any]:
|
||||
"""Anyrouter 使用 session_cookie 认证"""
|
||||
return AnyrouterConnector.get_credentials_schema()
|
||||
|
||||
@@ -394,9 +394,9 @@ class AnyrouterArchitecture(ProviderArchitecture):
|
||||
async def prepare_verify_config(
|
||||
self,
|
||||
base_url: str,
|
||||
config: Dict[str, Any],
|
||||
credentials: Dict[str, Any],
|
||||
) -> Dict[str, Any]:
|
||||
config: dict[str, Any],
|
||||
credentials: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
验证前获取 acw_sc__v2 Cookie
|
||||
|
||||
@@ -417,16 +417,16 @@ class AnyrouterArchitecture(ProviderArchitecture):
|
||||
|
||||
def build_verify_headers(
|
||||
self,
|
||||
config: Dict[str, Any],
|
||||
credentials: Dict[str, Any],
|
||||
) -> Dict[str, str]:
|
||||
config: dict[str, Any],
|
||||
credentials: dict[str, Any],
|
||||
) -> dict[str, str]:
|
||||
"""
|
||||
构建 Anyrouter 的验证请求 Headers
|
||||
|
||||
使用 Cookie 认证,不使用 Authorization。
|
||||
同时添加 New-Api-User header。
|
||||
"""
|
||||
headers: Dict[str, str] = {}
|
||||
headers: dict[str, str] = {}
|
||||
|
||||
cookies = []
|
||||
|
||||
@@ -455,7 +455,7 @@ class AnyrouterArchitecture(ProviderArchitecture):
|
||||
def parse_verify_response(
|
||||
self,
|
||||
status_code: int,
|
||||
data: Dict[str, Any],
|
||||
data: dict[str, Any],
|
||||
) -> VerifyResult:
|
||||
"""解析 Anyrouter 验证响应"""
|
||||
if status_code == 401:
|
||||
|
||||
@@ -6,7 +6,8 @@ from abc import ABC, abstractmethod
|
||||
from contextlib import asynccontextmanager
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, AsyncIterator, Dict, List, Optional, Type
|
||||
from typing import Any
|
||||
from collections.abc import AsyncIterator
|
||||
|
||||
import httpx
|
||||
|
||||
@@ -35,7 +36,7 @@ class ProviderConnector(ABC):
|
||||
auth_type: ConnectorAuthType = ConnectorAuthType.NONE
|
||||
display_name: str = "Base Connector"
|
||||
|
||||
def __init__(self, base_url: str, config: Optional[Dict[str, Any]] = None):
|
||||
def __init__(self, base_url: str, config: dict[str, Any] | None = None):
|
||||
"""
|
||||
初始化连接器
|
||||
|
||||
@@ -46,19 +47,19 @@ class ProviderConnector(ABC):
|
||||
self.base_url = base_url.rstrip("/")
|
||||
self.config = config or {}
|
||||
self._status = ConnectorStatus.DISCONNECTED
|
||||
self._connected_at: Optional[datetime] = None
|
||||
self._expires_at: Optional[datetime] = None
|
||||
self._last_error: Optional[str] = None
|
||||
self._connected_at: datetime | None = None
|
||||
self._expires_at: datetime | None = None
|
||||
self._last_error: str | None = None
|
||||
|
||||
# 代理配置
|
||||
self._proxy: Optional[str] = self.config.get("proxy")
|
||||
self._proxy: str | None = self.config.get("proxy")
|
||||
|
||||
# HTTP 客户端配置
|
||||
self._timeout = self.config.get("timeout", 30)
|
||||
self._headers: Dict[str, str] = {}
|
||||
self._headers: dict[str, str] = {}
|
||||
|
||||
@abstractmethod
|
||||
async def connect(self, credentials: Dict[str, Any]) -> bool:
|
||||
async def connect(self, credentials: dict[str, Any]) -> bool:
|
||||
"""
|
||||
建立认证连接
|
||||
|
||||
@@ -93,7 +94,7 @@ class ProviderConnector(ABC):
|
||||
"""
|
||||
pass
|
||||
|
||||
async def refresh_auth(self, credentials: Dict[str, Any]) -> bool:
|
||||
async def refresh_auth(self, credentials: dict[str, Any]) -> bool:
|
||||
"""
|
||||
刷新认证(如 Token 过期)
|
||||
|
||||
@@ -144,7 +145,7 @@ class ProviderConnector(ABC):
|
||||
last_error=self._last_error,
|
||||
)
|
||||
|
||||
def _set_connected(self, expires_at: Optional[datetime] = None) -> None:
|
||||
def _set_connected(self, expires_at: datetime | None = None) -> None:
|
||||
"""设置为已连接状态"""
|
||||
self._status = ConnectorStatus.CONNECTED
|
||||
self._connected_at = datetime.now(timezone.utc)
|
||||
@@ -163,7 +164,7 @@ class ProviderConnector(ABC):
|
||||
self._expires_at = None
|
||||
|
||||
@classmethod
|
||||
def get_credentials_schema(cls) -> Dict[str, Any]:
|
||||
def get_credentials_schema(cls) -> dict[str, Any]:
|
||||
"""
|
||||
获取凭据配置 JSON Schema(用于前端表单生成)
|
||||
|
||||
@@ -180,16 +181,16 @@ class VerifyResult:
|
||||
"""认证验证结果"""
|
||||
|
||||
success: bool
|
||||
message: Optional[str] = None
|
||||
username: Optional[str] = None
|
||||
display_name: Optional[str] = None
|
||||
email: Optional[str] = None
|
||||
quota: Optional[float] = None
|
||||
used_quota: Optional[float] = None
|
||||
request_count: Optional[int] = None
|
||||
extra: Optional[Dict[str, Any]] = None
|
||||
message: str | None = None
|
||||
username: str | None = None
|
||||
display_name: str | None = None
|
||||
email: str | None = None
|
||||
quota: float | None = None
|
||||
used_quota: float | None = None
|
||||
request_count: int | None = None
|
||||
extra: dict[str, Any] | None = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""转换为字典"""
|
||||
if not self.success:
|
||||
return {"success": False, "message": self.message}
|
||||
@@ -240,15 +241,15 @@ class ProviderArchitecture(ABC):
|
||||
description: str = ""
|
||||
|
||||
# 支持的 Connector 类型列表(按优先级排序)
|
||||
supported_connectors: List[Type[ProviderConnector]] = []
|
||||
supported_connectors: list[type[ProviderConnector]] = []
|
||||
|
||||
# 支持的 Action 类型列表
|
||||
supported_actions: List[Type[ProviderAction]] = []
|
||||
supported_actions: list[type[ProviderAction]] = []
|
||||
|
||||
# 默认操作配置
|
||||
default_action_configs: Dict[ProviderActionType, Dict[str, Any]] = {}
|
||||
default_action_configs: dict[ProviderActionType, dict[str, Any]] = {}
|
||||
|
||||
def __init__(self, config: Optional[Dict[str, Any]] = None):
|
||||
def __init__(self, config: dict[str, Any] | None = None):
|
||||
"""
|
||||
初始化架构
|
||||
|
||||
@@ -260,7 +261,7 @@ class ProviderArchitecture(ABC):
|
||||
# ==================== 认证验证相关方法(子类必须实现) ====================
|
||||
|
||||
@abstractmethod
|
||||
def get_credentials_schema(self) -> Dict[str, Any]:
|
||||
def get_credentials_schema(self) -> dict[str, Any]:
|
||||
"""
|
||||
获取凭据字段定义(JSON Schema 格式)
|
||||
|
||||
@@ -303,9 +304,9 @@ class ProviderArchitecture(ABC):
|
||||
@abstractmethod
|
||||
def build_verify_headers(
|
||||
self,
|
||||
config: Dict[str, Any],
|
||||
credentials: Dict[str, Any],
|
||||
) -> Dict[str, str]:
|
||||
config: dict[str, Any],
|
||||
credentials: dict[str, Any],
|
||||
) -> dict[str, str]:
|
||||
"""
|
||||
构建认证验证请求的 Headers
|
||||
|
||||
@@ -324,7 +325,7 @@ class ProviderArchitecture(ABC):
|
||||
def parse_verify_response(
|
||||
self,
|
||||
status_code: int,
|
||||
data: Dict[str, Any],
|
||||
data: dict[str, Any],
|
||||
) -> VerifyResult:
|
||||
"""
|
||||
解析认证验证响应
|
||||
@@ -345,9 +346,9 @@ class ProviderArchitecture(ABC):
|
||||
async def prepare_verify_config(
|
||||
self,
|
||||
base_url: str,
|
||||
config: Dict[str, Any],
|
||||
credentials: Dict[str, Any],
|
||||
) -> Dict[str, Any]:
|
||||
config: dict[str, Any],
|
||||
credentials: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
验证前的异步预处理(可选)
|
||||
|
||||
@@ -369,8 +370,8 @@ class ProviderArchitecture(ABC):
|
||||
def get_connector(
|
||||
self,
|
||||
base_url: str,
|
||||
auth_type: Optional[ConnectorAuthType] = None,
|
||||
config: Optional[Dict[str, Any]] = None,
|
||||
auth_type: ConnectorAuthType | None = None,
|
||||
config: dict[str, Any] | None = None,
|
||||
) -> ProviderConnector:
|
||||
"""
|
||||
获取连接器实例
|
||||
@@ -390,7 +391,7 @@ class ProviderArchitecture(ABC):
|
||||
raise ValueError(f"架构 {self.architecture_id} 未配置支持的连接器")
|
||||
|
||||
# 查找匹配的连接器
|
||||
connector_cls: Optional[Type[ProviderConnector]] = None
|
||||
connector_cls: type[ProviderConnector] | None = None
|
||||
|
||||
if auth_type:
|
||||
for cls in self.supported_connectors:
|
||||
@@ -413,7 +414,7 @@ class ProviderArchitecture(ABC):
|
||||
def get_action(
|
||||
self,
|
||||
action_type: ProviderActionType,
|
||||
config: Optional[Dict[str, Any]] = None,
|
||||
config: dict[str, Any] | None = None,
|
||||
) -> ProviderAction:
|
||||
"""
|
||||
获取操作实例
|
||||
@@ -428,7 +429,7 @@ class ProviderArchitecture(ABC):
|
||||
Raises:
|
||||
ValueError: 不支持的操作类型
|
||||
"""
|
||||
action_cls: Optional[Type[ProviderAction]] = None
|
||||
action_cls: type[ProviderAction] | None = None
|
||||
|
||||
for cls in self.supported_actions:
|
||||
if cls.action_type == action_type:
|
||||
@@ -457,15 +458,15 @@ class ProviderArchitecture(ABC):
|
||||
"""检查是否支持指定认证类型"""
|
||||
return any(c.auth_type == auth_type for c in self.supported_connectors)
|
||||
|
||||
def get_supported_auth_types(self) -> List[ConnectorAuthType]:
|
||||
def get_supported_auth_types(self) -> list[ConnectorAuthType]:
|
||||
"""获取支持的认证类型列表"""
|
||||
return [c.auth_type for c in self.supported_connectors]
|
||||
|
||||
def get_supported_action_types(self) -> List[ProviderActionType]:
|
||||
def get_supported_action_types(self) -> list[ProviderActionType]:
|
||||
"""获取支持的操作类型列表"""
|
||||
return [a.action_type for a in self.supported_actions]
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""转换为字典(用于 API 响应)"""
|
||||
return {
|
||||
"architecture_id": self.architecture_id,
|
||||
|
||||
@@ -4,7 +4,7 @@ Cubence 架构
|
||||
针对 Cubence 中转站的预设配置。
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, List, Optional, Type
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
@@ -54,11 +54,11 @@ class CubenceConnector(ProviderConnector):
|
||||
auth_type = ConnectorAuthType.COOKIE
|
||||
display_name = "Cubence Cookie"
|
||||
|
||||
def __init__(self, base_url: str, config: Optional[Dict[str, Any]] = None):
|
||||
def __init__(self, base_url: str, config: dict[str, Any] | None = None):
|
||||
super().__init__(base_url, config)
|
||||
self._token_cookie: Optional[str] = None
|
||||
self._token_cookie: str | None = None
|
||||
|
||||
async def connect(self, credentials: Dict[str, Any]) -> bool:
|
||||
async def connect(self, credentials: dict[str, Any]) -> bool:
|
||||
"""建立连接"""
|
||||
token_cookie = credentials.get("token_cookie")
|
||||
if not token_cookie:
|
||||
@@ -88,7 +88,7 @@ class CubenceConnector(ProviderConnector):
|
||||
return request
|
||||
|
||||
@classmethod
|
||||
def get_credentials_schema(cls) -> Dict[str, Any]:
|
||||
def get_credentials_schema(cls) -> dict[str, Any]:
|
||||
"""获取凭据配置 schema"""
|
||||
return {
|
||||
"type": "object",
|
||||
@@ -120,15 +120,15 @@ class CubenceArchitecture(ProviderArchitecture):
|
||||
display_name = "Cubence"
|
||||
description = "Cubence 中转站预设配置,使用 Cookie 认证"
|
||||
|
||||
supported_connectors: List[Type[ProviderConnector]] = [
|
||||
supported_connectors: list[type[ProviderConnector]] = [
|
||||
CubenceConnector,
|
||||
]
|
||||
|
||||
supported_actions: List[Type[ProviderAction]] = [
|
||||
supported_actions: list[type[ProviderAction]] = [
|
||||
CubenceBalanceAction,
|
||||
]
|
||||
|
||||
default_action_configs: Dict[ProviderActionType, Dict[str, Any]] = {
|
||||
default_action_configs: dict[ProviderActionType, dict[str, Any]] = {
|
||||
ProviderActionType.QUERY_BALANCE: {
|
||||
"endpoint": "/api/v1/dashboard/overview",
|
||||
"method": "GET",
|
||||
@@ -138,7 +138,7 @@ class CubenceArchitecture(ProviderArchitecture):
|
||||
},
|
||||
}
|
||||
|
||||
def get_credentials_schema(self) -> Dict[str, Any]:
|
||||
def get_credentials_schema(self) -> dict[str, Any]:
|
||||
"""Cubence 使用 token_cookie 认证"""
|
||||
return CubenceConnector.get_credentials_schema()
|
||||
|
||||
@@ -148,15 +148,15 @@ class CubenceArchitecture(ProviderArchitecture):
|
||||
|
||||
def build_verify_headers(
|
||||
self,
|
||||
config: Dict[str, Any],
|
||||
credentials: Dict[str, Any],
|
||||
) -> Dict[str, str]:
|
||||
config: dict[str, Any],
|
||||
credentials: dict[str, Any],
|
||||
) -> dict[str, str]:
|
||||
"""
|
||||
构建 Cubence 的验证请求 Headers
|
||||
|
||||
使用 Cookie 认证,不使用 Authorization。
|
||||
"""
|
||||
headers: Dict[str, str] = {}
|
||||
headers: dict[str, str] = {}
|
||||
|
||||
# 添加 token Cookie
|
||||
cookie_input = credentials.get("token_cookie")
|
||||
@@ -169,7 +169,7 @@ class CubenceArchitecture(ProviderArchitecture):
|
||||
def parse_verify_response(
|
||||
self,
|
||||
status_code: int,
|
||||
data: Dict[str, Any],
|
||||
data: dict[str, Any],
|
||||
) -> VerifyResult:
|
||||
"""解析 Cubence 验证响应"""
|
||||
if status_code == 401:
|
||||
@@ -190,7 +190,7 @@ class CubenceArchitecture(ProviderArchitecture):
|
||||
subscription_limits = user_data.get("subscription_limits", {})
|
||||
|
||||
# 构建 extra 信息,包含窗口限额
|
||||
extra: Dict[str, Any] = {
|
||||
extra: dict[str, Any] = {
|
||||
"role": user_info.get("role"),
|
||||
"invite_code": user_info.get("invite_code"),
|
||||
}
|
||||
|
||||
@@ -39,7 +39,7 @@
|
||||
3. 在前端 auth-templates/ 添加对应的模板定义
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, List, Optional, Type
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
@@ -67,14 +67,14 @@ class GenericApiKeyConnector(ProviderConnector):
|
||||
auth_type = ConnectorAuthType.API_KEY
|
||||
display_name = "API Key"
|
||||
|
||||
def __init__(self, base_url: str, config: Optional[Dict[str, Any]] = None):
|
||||
def __init__(self, base_url: str, config: dict[str, Any] | None = None):
|
||||
super().__init__(base_url, config)
|
||||
self._api_key: Optional[str] = None
|
||||
self._api_key: str | None = None
|
||||
# 支持配置认证方式
|
||||
self._auth_method = self.config.get("auth_method", "bearer")
|
||||
self._header_name = self.config.get("header_name", "Authorization")
|
||||
|
||||
async def connect(self, credentials: Dict[str, Any]) -> bool:
|
||||
async def connect(self, credentials: dict[str, Any]) -> bool:
|
||||
"""建立连接"""
|
||||
api_key = credentials.get("api_key")
|
||||
if not api_key:
|
||||
@@ -107,7 +107,7 @@ class GenericApiKeyConnector(ProviderConnector):
|
||||
return request
|
||||
|
||||
@classmethod
|
||||
def get_credentials_schema(cls) -> Dict[str, Any]:
|
||||
def get_credentials_schema(cls) -> dict[str, Any]:
|
||||
"""获取凭据配置 schema"""
|
||||
return {
|
||||
"type": "object",
|
||||
@@ -136,14 +136,14 @@ class GenericApiArchitecture(ProviderArchitecture):
|
||||
display_name = "通用 API"
|
||||
description = "可配置的通用 API 架构,适用于各种中转站"
|
||||
|
||||
supported_connectors: List[Type[ProviderConnector]] = [
|
||||
supported_connectors: list[type[ProviderConnector]] = [
|
||||
GenericApiKeyConnector,
|
||||
]
|
||||
|
||||
supported_actions: List[Type[ProviderAction]] = [NewApiBalanceAction]
|
||||
supported_actions: list[type[ProviderAction]] = [NewApiBalanceAction]
|
||||
|
||||
# 默认操作配置(可被用户配置覆盖)
|
||||
default_action_configs: Dict[ProviderActionType, Dict[str, Any]] = {
|
||||
default_action_configs: dict[ProviderActionType, dict[str, Any]] = {
|
||||
ProviderActionType.QUERY_BALANCE: {
|
||||
"endpoint": "/api/user/balance",
|
||||
"method": "GET",
|
||||
@@ -154,7 +154,7 @@ class GenericApiArchitecture(ProviderArchitecture):
|
||||
},
|
||||
}
|
||||
|
||||
def get_credentials_schema(self) -> Dict[str, Any]:
|
||||
def get_credentials_schema(self) -> dict[str, Any]:
|
||||
"""通用架构只需要 api_key"""
|
||||
return GenericApiKeyConnector.get_credentials_schema()
|
||||
|
||||
@@ -164,11 +164,11 @@ class GenericApiArchitecture(ProviderArchitecture):
|
||||
|
||||
def build_verify_headers(
|
||||
self,
|
||||
config: Dict[str, Any],
|
||||
credentials: Dict[str, Any],
|
||||
) -> Dict[str, str]:
|
||||
config: dict[str, Any],
|
||||
credentials: dict[str, Any],
|
||||
) -> dict[str, str]:
|
||||
"""构建通用 API 的验证请求 Headers"""
|
||||
headers: Dict[str, str] = {}
|
||||
headers: dict[str, str] = {}
|
||||
|
||||
api_key = credentials.get("api_key", "")
|
||||
if api_key:
|
||||
@@ -184,7 +184,7 @@ class GenericApiArchitecture(ProviderArchitecture):
|
||||
def parse_verify_response(
|
||||
self,
|
||||
status_code: int,
|
||||
data: Dict[str, Any],
|
||||
data: dict[str, Any],
|
||||
) -> VerifyResult:
|
||||
"""解析通用 API 验证响应"""
|
||||
if status_code == 401:
|
||||
|
||||
@@ -4,7 +4,7 @@ NekoCode 架构
|
||||
针对 NekoCode 中转站的预设配置,使用 Cookie 认证。
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, List, Optional, Type
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
@@ -56,11 +56,11 @@ class NekoCodeConnector(ProviderConnector):
|
||||
auth_type = ConnectorAuthType.COOKIE
|
||||
display_name = "NekoCode Cookie"
|
||||
|
||||
def __init__(self, base_url: str, config: Optional[Dict[str, Any]] = None):
|
||||
def __init__(self, base_url: str, config: dict[str, Any] | None = None):
|
||||
super().__init__(base_url, config)
|
||||
self._session_cookie: Optional[str] = None
|
||||
self._session_cookie: str | None = None
|
||||
|
||||
async def connect(self, credentials: Dict[str, Any]) -> bool:
|
||||
async def connect(self, credentials: dict[str, Any]) -> bool:
|
||||
"""建立连接"""
|
||||
session_cookie = credentials.get("session_cookie")
|
||||
if not session_cookie:
|
||||
@@ -90,7 +90,7 @@ class NekoCodeConnector(ProviderConnector):
|
||||
return request
|
||||
|
||||
@classmethod
|
||||
def get_credentials_schema(cls) -> Dict[str, Any]:
|
||||
def get_credentials_schema(cls) -> dict[str, Any]:
|
||||
"""获取凭据配置 schema"""
|
||||
return {
|
||||
"type": "object",
|
||||
@@ -121,20 +121,20 @@ class NekoCodeArchitecture(ProviderArchitecture):
|
||||
display_name = "NekoCode"
|
||||
description = "NekoCode 中转站预设配置,使用 Cookie 认证"
|
||||
|
||||
supported_connectors: List[Type[ProviderConnector]] = [
|
||||
supported_connectors: list[type[ProviderConnector]] = [
|
||||
NekoCodeConnector,
|
||||
]
|
||||
|
||||
supported_actions: List[Type[ProviderAction]] = [NekoCodeBalanceAction]
|
||||
supported_actions: list[type[ProviderAction]] = [NekoCodeBalanceAction]
|
||||
|
||||
default_action_configs: Dict[ProviderActionType, Dict[str, Any]] = {
|
||||
default_action_configs: dict[ProviderActionType, dict[str, Any]] = {
|
||||
ProviderActionType.QUERY_BALANCE: {
|
||||
"endpoint": "/api/usage/summary",
|
||||
"method": "GET",
|
||||
},
|
||||
}
|
||||
|
||||
def get_credentials_schema(self) -> Dict[str, Any]:
|
||||
def get_credentials_schema(self) -> dict[str, Any]:
|
||||
"""NekoCode 使用 session_cookie 认证"""
|
||||
return NekoCodeConnector.get_credentials_schema()
|
||||
|
||||
@@ -145,9 +145,9 @@ class NekoCodeArchitecture(ProviderArchitecture):
|
||||
async def prepare_verify_config(
|
||||
self,
|
||||
base_url: str,
|
||||
config: Dict[str, Any],
|
||||
credentials: Dict[str, Any],
|
||||
) -> Dict[str, Any]:
|
||||
config: dict[str, Any],
|
||||
credentials: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
验证前获取 /api/usage/summary 数据(用于显示天卡信息)
|
||||
|
||||
@@ -161,7 +161,7 @@ class NekoCodeArchitecture(ProviderArchitecture):
|
||||
"""
|
||||
try:
|
||||
# 构建请求头
|
||||
headers: Dict[str, str] = {
|
||||
headers: dict[str, str] = {
|
||||
"User-Agent": (
|
||||
"Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
|
||||
"AppleWebKit/537.36 (KHTML, like Gecko) "
|
||||
@@ -176,7 +176,7 @@ class NekoCodeArchitecture(ProviderArchitecture):
|
||||
headers["Cookie"] = f"session={session_value}"
|
||||
|
||||
# 构建 client 参数
|
||||
client_kwargs: Dict[str, Any] = {
|
||||
client_kwargs: dict[str, Any] = {
|
||||
"timeout": 10,
|
||||
"verify": get_ssl_context(),
|
||||
}
|
||||
@@ -202,15 +202,15 @@ class NekoCodeArchitecture(ProviderArchitecture):
|
||||
|
||||
def build_verify_headers(
|
||||
self,
|
||||
config: Dict[str, Any],
|
||||
credentials: Dict[str, Any],
|
||||
) -> Dict[str, str]:
|
||||
config: dict[str, Any],
|
||||
credentials: dict[str, Any],
|
||||
) -> dict[str, str]:
|
||||
"""
|
||||
构建 NekoCode 的验证请求 Headers
|
||||
|
||||
使用 Cookie 认证,不使用 Authorization。
|
||||
"""
|
||||
headers: Dict[str, str] = {}
|
||||
headers: dict[str, str] = {}
|
||||
|
||||
# 添加 session Cookie
|
||||
cookie_input = credentials.get("session_cookie")
|
||||
@@ -224,7 +224,7 @@ class NekoCodeArchitecture(ProviderArchitecture):
|
||||
def parse_verify_response(
|
||||
self,
|
||||
status_code: int,
|
||||
data: Dict[str, Any],
|
||||
data: dict[str, Any],
|
||||
) -> VerifyResult:
|
||||
"""解析 NekoCode 验证响应(/api/user/self + _usage_summary)"""
|
||||
if status_code == 401:
|
||||
@@ -249,7 +249,7 @@ class NekoCodeArchitecture(ProviderArchitecture):
|
||||
quota = None
|
||||
|
||||
# 从 prepare_verify_config 获取的 _usage_summary 数据(天卡信息)
|
||||
extra: Dict[str, Any] = {}
|
||||
extra: dict[str, Any] = {}
|
||||
usage_summary = data.get("_usage_summary", {})
|
||||
subscription = usage_summary.get("subscription", {})
|
||||
|
||||
|
||||
@@ -4,7 +4,7 @@ New API 架构
|
||||
针对 New API 风格的中转站优化的预设配置。
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, List, Optional, Type
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
@@ -32,13 +32,13 @@ class NewApiConnector(ProviderConnector):
|
||||
auth_type = ConnectorAuthType.API_KEY
|
||||
display_name = "New API Key"
|
||||
|
||||
def __init__(self, base_url: str, config: Optional[Dict[str, Any]] = None):
|
||||
def __init__(self, base_url: str, config: dict[str, Any] | None = None):
|
||||
super().__init__(base_url, config)
|
||||
self._api_key: Optional[str] = None
|
||||
self._user_id: Optional[str] = None
|
||||
self._cookie: Optional[str] = None
|
||||
self._api_key: str | None = None
|
||||
self._user_id: str | None = None
|
||||
self._cookie: str | None = None
|
||||
|
||||
async def connect(self, credentials: Dict[str, Any]) -> bool:
|
||||
async def connect(self, credentials: dict[str, Any]) -> bool:
|
||||
"""建立连接"""
|
||||
api_key = credentials.get("api_key")
|
||||
cookie = credentials.get("cookie")
|
||||
@@ -85,7 +85,7 @@ class NewApiConnector(ProviderConnector):
|
||||
return request
|
||||
|
||||
@classmethod
|
||||
def get_credentials_schema(cls) -> Dict[str, Any]:
|
||||
def get_credentials_schema(cls) -> dict[str, Any]:
|
||||
"""获取凭据配置 schema"""
|
||||
return {
|
||||
"type": "object",
|
||||
@@ -127,15 +127,15 @@ class NewApiArchitecture(ProviderArchitecture):
|
||||
display_name = "New API"
|
||||
description = "New API 风格中转站的预设配置"
|
||||
|
||||
supported_connectors: List[Type[ProviderConnector]] = [
|
||||
supported_connectors: list[type[ProviderConnector]] = [
|
||||
NewApiConnector,
|
||||
]
|
||||
|
||||
supported_actions: List[Type[ProviderAction]] = [
|
||||
supported_actions: list[type[ProviderAction]] = [
|
||||
NewApiBalanceAction,
|
||||
]
|
||||
|
||||
default_action_configs: Dict[ProviderActionType, Dict[str, Any]] = {
|
||||
default_action_configs: dict[ProviderActionType, dict[str, Any]] = {
|
||||
ProviderActionType.QUERY_BALANCE: {
|
||||
"endpoint": "/api/user/self",
|
||||
"method": "GET",
|
||||
@@ -144,7 +144,7 @@ class NewApiArchitecture(ProviderArchitecture):
|
||||
},
|
||||
}
|
||||
|
||||
def get_credentials_schema(self) -> Dict[str, Any]:
|
||||
def get_credentials_schema(self) -> dict[str, Any]:
|
||||
"""New API 需要 api_key 和 user_id"""
|
||||
return NewApiConnector.get_credentials_schema()
|
||||
|
||||
@@ -154,15 +154,15 @@ class NewApiArchitecture(ProviderArchitecture):
|
||||
|
||||
def build_verify_headers(
|
||||
self,
|
||||
config: Dict[str, Any],
|
||||
credentials: Dict[str, Any],
|
||||
) -> Dict[str, str]:
|
||||
config: dict[str, Any],
|
||||
credentials: dict[str, Any],
|
||||
) -> dict[str, str]:
|
||||
"""
|
||||
构建 New API 的验证请求 Headers
|
||||
|
||||
New API 特有:需要 New-Api-User Header 传递用户 ID
|
||||
"""
|
||||
headers: Dict[str, str] = {}
|
||||
headers: dict[str, str] = {}
|
||||
|
||||
# Bearer Token 认证
|
||||
api_key = credentials.get("api_key", "")
|
||||
@@ -184,7 +184,7 @@ class NewApiArchitecture(ProviderArchitecture):
|
||||
def parse_verify_response(
|
||||
self,
|
||||
status_code: int,
|
||||
data: Dict[str, Any],
|
||||
data: dict[str, Any],
|
||||
) -> VerifyResult:
|
||||
"""解析 New API 验证响应"""
|
||||
if status_code == 401:
|
||||
|
||||
@@ -4,7 +4,7 @@ One API 架构
|
||||
针对 One API 风格的中转站优化的预设配置。
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, List, Optional, Type
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
@@ -29,11 +29,11 @@ class OneApiConnector(ProviderConnector):
|
||||
auth_type = ConnectorAuthType.API_KEY
|
||||
display_name = "One API Key"
|
||||
|
||||
def __init__(self, base_url: str, config: Optional[Dict[str, Any]] = None):
|
||||
def __init__(self, base_url: str, config: dict[str, Any] | None = None):
|
||||
super().__init__(base_url, config)
|
||||
self._api_key: Optional[str] = None
|
||||
self._api_key: str | None = None
|
||||
|
||||
async def connect(self, credentials: Dict[str, Any]) -> bool:
|
||||
async def connect(self, credentials: dict[str, Any]) -> bool:
|
||||
"""建立连接"""
|
||||
api_key = credentials.get("api_key")
|
||||
if not api_key:
|
||||
@@ -60,7 +60,7 @@ class OneApiConnector(ProviderConnector):
|
||||
return request
|
||||
|
||||
@classmethod
|
||||
def get_credentials_schema(cls) -> Dict[str, Any]:
|
||||
def get_credentials_schema(cls) -> dict[str, Any]:
|
||||
"""获取凭据配置 schema"""
|
||||
return {
|
||||
"type": "object",
|
||||
@@ -91,15 +91,15 @@ class OneApiArchitecture(ProviderArchitecture):
|
||||
display_name = "One API"
|
||||
description = "One API 风格中转站的预设配置"
|
||||
|
||||
supported_connectors: List[Type[ProviderConnector]] = [
|
||||
supported_connectors: list[type[ProviderConnector]] = [
|
||||
OneApiConnector,
|
||||
]
|
||||
|
||||
supported_actions: List[Type[ProviderAction]] = [
|
||||
supported_actions: list[type[ProviderAction]] = [
|
||||
NewApiBalanceAction,
|
||||
]
|
||||
|
||||
default_action_configs: Dict[ProviderActionType, Dict[str, Any]] = {
|
||||
default_action_configs: dict[ProviderActionType, dict[str, Any]] = {
|
||||
ProviderActionType.QUERY_BALANCE: {
|
||||
"endpoint": "/api/user/self",
|
||||
"method": "GET",
|
||||
@@ -110,7 +110,7 @@ class OneApiArchitecture(ProviderArchitecture):
|
||||
},
|
||||
}
|
||||
|
||||
def get_credentials_schema(self) -> Dict[str, Any]:
|
||||
def get_credentials_schema(self) -> dict[str, Any]:
|
||||
"""One API 只需要 api_key"""
|
||||
return OneApiConnector.get_credentials_schema()
|
||||
|
||||
@@ -120,11 +120,11 @@ class OneApiArchitecture(ProviderArchitecture):
|
||||
|
||||
def build_verify_headers(
|
||||
self,
|
||||
config: Dict[str, Any],
|
||||
credentials: Dict[str, Any],
|
||||
) -> Dict[str, str]:
|
||||
config: dict[str, Any],
|
||||
credentials: dict[str, Any],
|
||||
) -> dict[str, str]:
|
||||
"""构建 One API 的验证请求 Headers"""
|
||||
headers: Dict[str, str] = {}
|
||||
headers: dict[str, str] = {}
|
||||
|
||||
api_key = credentials.get("api_key", "")
|
||||
if api_key:
|
||||
@@ -135,7 +135,7 @@ class OneApiArchitecture(ProviderArchitecture):
|
||||
def parse_verify_response(
|
||||
self,
|
||||
status_code: int,
|
||||
data: Dict[str, Any],
|
||||
data: dict[str, Any],
|
||||
) -> VerifyResult:
|
||||
"""解析 One API 验证响应"""
|
||||
if status_code == 401:
|
||||
|
||||
@@ -4,7 +4,7 @@ YesCode 架构
|
||||
针对 YesCode 中转站的预设配置。
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, List, Optional, Type
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
@@ -23,7 +23,7 @@ from src.services.provider_ops.types import ConnectorAuthType, ProviderActionTyp
|
||||
from src.utils.ssl_utils import get_ssl_context
|
||||
|
||||
|
||||
def _extract_cookies(cookie_string: str) -> Dict[str, str]:
|
||||
def _extract_cookies(cookie_string: str) -> dict[str, str]:
|
||||
"""
|
||||
从完整的 Cookie 字符串中提取 yescode_auth 和 yescode_csrf
|
||||
|
||||
@@ -33,7 +33,7 @@ def _extract_cookies(cookie_string: str) -> Dict[str, str]:
|
||||
Returns:
|
||||
包含 yescode_auth 和 yescode_csrf 的字典
|
||||
"""
|
||||
result: Dict[str, str] = {}
|
||||
result: dict[str, str] = {}
|
||||
for part in cookie_string.split(";"):
|
||||
part = part.strip()
|
||||
if "=" in part:
|
||||
@@ -82,11 +82,11 @@ class YesCodeConnector(ProviderConnector):
|
||||
auth_type = ConnectorAuthType.COOKIE
|
||||
display_name = "YesCode Cookie"
|
||||
|
||||
def __init__(self, base_url: str, config: Optional[Dict[str, Any]] = None):
|
||||
def __init__(self, base_url: str, config: dict[str, Any] | None = None):
|
||||
super().__init__(base_url, config)
|
||||
self._auth_cookie: Optional[str] = None
|
||||
self._auth_cookie: str | None = None
|
||||
|
||||
async def connect(self, credentials: Dict[str, Any]) -> bool:
|
||||
async def connect(self, credentials: dict[str, Any]) -> bool:
|
||||
"""建立连接"""
|
||||
auth_cookie = credentials.get("auth_cookie")
|
||||
if not auth_cookie:
|
||||
@@ -116,7 +116,7 @@ class YesCodeConnector(ProviderConnector):
|
||||
return request
|
||||
|
||||
@classmethod
|
||||
def get_credentials_schema(cls) -> Dict[str, Any]:
|
||||
def get_credentials_schema(cls) -> dict[str, Any]:
|
||||
"""获取凭据配置 schema"""
|
||||
return {
|
||||
"type": "object",
|
||||
@@ -148,15 +148,15 @@ class YesCodeArchitecture(ProviderArchitecture):
|
||||
display_name = "YesCode"
|
||||
description = "YesCode 中转站预设配置,使用 Cookie 认证"
|
||||
|
||||
supported_connectors: List[Type[ProviderConnector]] = [
|
||||
supported_connectors: list[type[ProviderConnector]] = [
|
||||
YesCodeConnector,
|
||||
]
|
||||
|
||||
supported_actions: List[Type[ProviderAction]] = [
|
||||
supported_actions: list[type[ProviderAction]] = [
|
||||
YesCodeBalanceAction,
|
||||
]
|
||||
|
||||
default_action_configs: Dict[ProviderActionType, Dict[str, Any]] = {
|
||||
default_action_configs: dict[ProviderActionType, dict[str, Any]] = {
|
||||
ProviderActionType.QUERY_BALANCE: {
|
||||
"endpoint": "/api/v1/user/balance",
|
||||
"method": "GET",
|
||||
@@ -166,7 +166,7 @@ class YesCodeArchitecture(ProviderArchitecture):
|
||||
},
|
||||
}
|
||||
|
||||
def get_credentials_schema(self) -> Dict[str, Any]:
|
||||
def get_credentials_schema(self) -> dict[str, Any]:
|
||||
"""YesCode 使用 auth_cookie 认证"""
|
||||
return YesCodeConnector.get_credentials_schema()
|
||||
|
||||
@@ -177,15 +177,15 @@ class YesCodeArchitecture(ProviderArchitecture):
|
||||
async def prepare_verify_config(
|
||||
self,
|
||||
base_url: str,
|
||||
config: Dict[str, Any],
|
||||
credentials: Dict[str, Any],
|
||||
) -> Dict[str, Any]:
|
||||
config: dict[str, Any],
|
||||
credentials: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
预获取合并数据(balance + profile)
|
||||
|
||||
验证时并发调用两个接口获取完整数据。
|
||||
"""
|
||||
extra_config: Dict[str, Any] = {}
|
||||
extra_config: dict[str, Any] = {}
|
||||
|
||||
cookie_input = credentials.get("auth_cookie")
|
||||
if not cookie_input:
|
||||
@@ -198,7 +198,7 @@ class YesCodeArchitecture(ProviderArchitecture):
|
||||
|
||||
try:
|
||||
# 构建 client 参数
|
||||
client_kwargs: Dict[str, Any] = {
|
||||
client_kwargs: dict[str, Any] = {
|
||||
"headers": {"Cookie": cookie_header},
|
||||
"timeout": 10.0,
|
||||
"verify": get_ssl_context(),
|
||||
@@ -218,15 +218,15 @@ class YesCodeArchitecture(ProviderArchitecture):
|
||||
|
||||
def build_verify_headers(
|
||||
self,
|
||||
config: Dict[str, Any],
|
||||
credentials: Dict[str, Any],
|
||||
) -> Dict[str, str]:
|
||||
config: dict[str, Any],
|
||||
credentials: dict[str, Any],
|
||||
) -> dict[str, str]:
|
||||
"""
|
||||
构建 YesCode 的验证请求 Headers
|
||||
|
||||
使用 Cookie 认证,不使用 Authorization。
|
||||
"""
|
||||
headers: Dict[str, str] = {}
|
||||
headers: dict[str, str] = {}
|
||||
|
||||
# 添加 Cookie
|
||||
cookie_input = credentials.get("auth_cookie")
|
||||
@@ -238,7 +238,7 @@ class YesCodeArchitecture(ProviderArchitecture):
|
||||
def parse_verify_response(
|
||||
self,
|
||||
status_code: int,
|
||||
data: Dict[str, Any],
|
||||
data: dict[str, Any],
|
||||
) -> VerifyResult:
|
||||
"""解析 YesCode 验证响应(使用预获取的合并数据)"""
|
||||
if status_code == 401:
|
||||
|
||||
@@ -4,8 +4,9 @@
|
||||
管理所有可用的 Provider 架构。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
from typing import Dict, List, Optional, Type
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.services.provider_ops.architectures import (
|
||||
@@ -27,11 +28,11 @@ class ArchitectureRegistry:
|
||||
单例模式,管理所有可用的 Provider 架构。
|
||||
"""
|
||||
|
||||
_instance: Optional["ArchitectureRegistry"] = None
|
||||
_instance: ArchitectureRegistry | None = None
|
||||
_lock = threading.Lock()
|
||||
_initialized: bool = False
|
||||
|
||||
def __new__(cls) -> "ArchitectureRegistry":
|
||||
def __new__(cls) -> ArchitectureRegistry:
|
||||
if cls._instance is None:
|
||||
with cls._lock:
|
||||
if cls._instance is None:
|
||||
@@ -43,7 +44,7 @@ class ArchitectureRegistry:
|
||||
if self._initialized:
|
||||
return
|
||||
|
||||
self._architectures: Dict[str, ProviderArchitecture] = {}
|
||||
self._architectures: dict[str, ProviderArchitecture] = {}
|
||||
self._initialized = True
|
||||
|
||||
# 注册内置架构
|
||||
@@ -51,7 +52,7 @@ class ArchitectureRegistry:
|
||||
|
||||
def _register_builtin_architectures(self) -> None:
|
||||
"""注册内置架构"""
|
||||
builtin: List[Type[ProviderArchitecture]] = [
|
||||
builtin: list[type[ProviderArchitecture]] = [
|
||||
AnyrouterArchitecture,
|
||||
CubenceArchitecture,
|
||||
GenericApiArchitecture,
|
||||
@@ -92,7 +93,7 @@ class ArchitectureRegistry:
|
||||
return True
|
||||
return False
|
||||
|
||||
def get(self, architecture_id: str) -> Optional[ProviderArchitecture]:
|
||||
def get(self, architecture_id: str) -> ProviderArchitecture | None:
|
||||
"""
|
||||
获取架构
|
||||
|
||||
@@ -104,7 +105,7 @@ class ArchitectureRegistry:
|
||||
"""
|
||||
return self._architectures.get(architecture_id)
|
||||
|
||||
def get_or_default(self, architecture_id: Optional[str] = None) -> ProviderArchitecture:
|
||||
def get_or_default(self, architecture_id: str | None = None) -> ProviderArchitecture:
|
||||
"""
|
||||
获取架构,如果不存在则返回默认架构
|
||||
|
||||
@@ -120,21 +121,21 @@ class ArchitectureRegistry:
|
||||
# 返回默认架构(generic_api)
|
||||
return self._architectures.get("generic_api", GenericApiArchitecture())
|
||||
|
||||
def list_all(self) -> List[ProviderArchitecture]:
|
||||
def list_all(self) -> list[ProviderArchitecture]:
|
||||
"""获取所有已注册的架构"""
|
||||
return list(self._architectures.values())
|
||||
|
||||
def list_ids(self) -> List[str]:
|
||||
def list_ids(self) -> list[str]:
|
||||
"""获取所有已注册的架构 ID"""
|
||||
return list(self._architectures.keys())
|
||||
|
||||
def to_dict_list(self) -> List[Dict]:
|
||||
def to_dict_list(self) -> list[dict]:
|
||||
"""获取所有架构的字典表示(用于 API 响应)"""
|
||||
return [arch.to_dict() for arch in self._architectures.values()]
|
||||
|
||||
|
||||
# 全局注册表实例
|
||||
_registry: Optional[ArchitectureRegistry] = None
|
||||
_registry: ArchitectureRegistry | None = None
|
||||
|
||||
|
||||
def get_registry() -> ArchitectureRegistry:
|
||||
|
||||
@@ -8,7 +8,7 @@ import asyncio
|
||||
import os
|
||||
from dataclasses import asdict
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Dict, List, Optional
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
@@ -18,7 +18,7 @@ from src.core.crypto import CryptoService
|
||||
from src.core.logger import logger
|
||||
from src.database import create_session
|
||||
from src.models.database import Provider
|
||||
from src.services.provider_ops.architectures import ProviderArchitecture, ProviderConnector
|
||||
from src.services.provider_ops.architectures import ProviderConnector
|
||||
from src.services.provider_ops.registry import get_registry
|
||||
from src.services.provider_ops.types import (
|
||||
ActionResult,
|
||||
@@ -85,11 +85,11 @@ class ProviderOpsService:
|
||||
self.crypto = CryptoService()
|
||||
|
||||
# 连接器缓存 {provider_id: ProviderConnector}
|
||||
self._connectors: Dict[str, ProviderConnector] = {}
|
||||
self._connectors: dict[str, ProviderConnector] = {}
|
||||
|
||||
# ==================== 配置管理 ====================
|
||||
|
||||
def get_config(self, provider_id: str) -> Optional[ProviderOpsConfig]:
|
||||
def get_config(self, provider_id: str) -> ProviderOpsConfig | None:
|
||||
"""
|
||||
获取 Provider 的操作配置
|
||||
|
||||
@@ -186,7 +186,7 @@ class ProviderOpsService:
|
||||
async def connect(
|
||||
self,
|
||||
provider_id: str,
|
||||
credentials: Optional[Dict[str, Any]] = None,
|
||||
credentials: dict[str, Any] | None = None,
|
||||
) -> tuple[bool, str]:
|
||||
"""
|
||||
建立与 Provider 的连接
|
||||
@@ -296,7 +296,7 @@ class ProviderOpsService:
|
||||
self,
|
||||
provider_id: str,
|
||||
action_type: ProviderActionType,
|
||||
action_config: Optional[Dict[str, Any]] = None,
|
||||
action_config: dict[str, Any] | None = None,
|
||||
) -> ActionResult:
|
||||
"""
|
||||
执行操作
|
||||
@@ -371,7 +371,7 @@ class ProviderOpsService:
|
||||
async def query_balance(
|
||||
self,
|
||||
provider_id: str,
|
||||
config: Optional[Dict[str, Any]] = None,
|
||||
config: dict[str, Any] | None = None,
|
||||
) -> ActionResult:
|
||||
"""
|
||||
查询余额(快捷方法)
|
||||
@@ -508,7 +508,7 @@ class ProviderOpsService:
|
||||
self,
|
||||
provider_id: str,
|
||||
quota_usd: float,
|
||||
extra: Optional[Dict[str, Any]] = None,
|
||||
extra: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
从验证结果缓存余额
|
||||
@@ -538,7 +538,7 @@ class ProviderOpsService:
|
||||
await CacheService.set(cache_key, cache_data, BALANCE_CACHE_TTL)
|
||||
logger.debug(f"验证成功,缓存余额: provider_id={provider_id}, quota_usd={quota_usd}")
|
||||
|
||||
async def _get_cached_balance(self, provider_id: str) -> Optional[ActionResult]:
|
||||
async def _get_cached_balance(self, provider_id: str) -> ActionResult | None:
|
||||
"""获取缓存的余额"""
|
||||
cache_key = f"provider_ops:balance:{provider_id}"
|
||||
cached = await CacheService.get(cache_key)
|
||||
@@ -585,7 +585,7 @@ class ProviderOpsService:
|
||||
async def checkin(
|
||||
self,
|
||||
provider_id: str,
|
||||
config: Optional[Dict[str, Any]] = None,
|
||||
config: dict[str, Any] | None = None,
|
||||
) -> ActionResult:
|
||||
"""
|
||||
签到(快捷方法)
|
||||
@@ -601,11 +601,11 @@ class ProviderOpsService:
|
||||
|
||||
# ==================== 辅助方法 ====================
|
||||
|
||||
def _get_provider(self, provider_id: str) -> Optional[Provider]:
|
||||
def _get_provider(self, provider_id: str) -> Provider | None:
|
||||
"""获取 Provider"""
|
||||
return self.db.query(Provider).filter(Provider.id == provider_id).first()
|
||||
|
||||
def _get_provider_base_url(self, provider: Provider) -> Optional[str]:
|
||||
def _get_provider_base_url(self, provider: Provider) -> str | None:
|
||||
"""从 Provider 获取 base_url"""
|
||||
# 优先从第一个 endpoint 获取
|
||||
if provider.endpoints:
|
||||
@@ -624,7 +624,7 @@ class ProviderOpsService:
|
||||
|
||||
return None
|
||||
|
||||
def _encrypt_credentials(self, credentials: Dict[str, Any]) -> Dict[str, Any]:
|
||||
def _encrypt_credentials(self, credentials: dict[str, Any]) -> dict[str, Any]:
|
||||
"""加密凭据中的敏感字段"""
|
||||
encrypted = {}
|
||||
for key, value in credentials.items():
|
||||
@@ -639,7 +639,7 @@ class ProviderOpsService:
|
||||
encrypted[key] = value
|
||||
return encrypted
|
||||
|
||||
def _decrypt_credentials(self, credentials: Dict[str, Any]) -> Dict[str, Any]:
|
||||
def _decrypt_credentials(self, credentials: dict[str, Any]) -> dict[str, Any]:
|
||||
"""解密凭据中的敏感字段"""
|
||||
decrypted = {}
|
||||
for key, value in credentials.items():
|
||||
@@ -653,7 +653,7 @@ class ProviderOpsService:
|
||||
decrypted[key] = value
|
||||
return decrypted
|
||||
|
||||
def get_masked_credentials(self, credentials: Dict[str, Any]) -> Dict[str, Any]:
|
||||
def get_masked_credentials(self, credentials: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
获取脱敏后的凭据
|
||||
|
||||
@@ -683,8 +683,8 @@ class ProviderOpsService:
|
||||
def merge_credentials_with_saved(
|
||||
self,
|
||||
provider_id: str,
|
||||
credentials: Dict[str, Any],
|
||||
) -> Dict[str, Any]:
|
||||
credentials: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
合并凭据:如果请求中的敏感字段为空,使用已保存的凭据
|
||||
|
||||
@@ -720,8 +720,8 @@ class ProviderOpsService:
|
||||
# ==================== 批量操作 ====================
|
||||
|
||||
async def batch_query_balance(
|
||||
self, provider_ids: Optional[List[str]] = None
|
||||
) -> Dict[str, ActionResult]:
|
||||
self, provider_ids: list[str] | None = None
|
||||
) -> dict[str, ActionResult]:
|
||||
"""
|
||||
批量查询余额(优先返回缓存,后台异步刷新)
|
||||
|
||||
@@ -780,10 +780,10 @@ class ProviderOpsService:
|
||||
base_url: str,
|
||||
architecture_id: str,
|
||||
auth_type: ConnectorAuthType,
|
||||
config: Dict[str, Any],
|
||||
credentials: Dict[str, Any],
|
||||
provider_id: Optional[str] = None,
|
||||
) -> Dict[str, Any]:
|
||||
config: dict[str, Any],
|
||||
credentials: dict[str, Any],
|
||||
provider_id: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
验证认证配置
|
||||
|
||||
@@ -831,7 +831,7 @@ class ProviderOpsService:
|
||||
|
||||
try:
|
||||
# 构建 httpx client 参数
|
||||
client_kwargs: Dict[str, Any] = {
|
||||
client_kwargs: dict[str, Any] = {
|
||||
"timeout": 30.0,
|
||||
"verify": get_ssl_context(),
|
||||
}
|
||||
|
||||
@@ -2,10 +2,12 @@
|
||||
Provider 操作模块类型定义
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
from enum import Enum
|
||||
from typing import Any, Dict, List, Optional
|
||||
from typing import Any
|
||||
|
||||
|
||||
class ConnectorAuthType(str, Enum):
|
||||
@@ -60,23 +62,23 @@ class ConnectorStatus(str, Enum):
|
||||
class BalanceInfo:
|
||||
"""余额信息"""
|
||||
|
||||
total_granted: Optional[float] = None # 总授予额度
|
||||
total_used: Optional[float] = None # 已使用额度
|
||||
total_available: Optional[float] = None # 可用余额
|
||||
expires_at: Optional[datetime] = None # 过期时间
|
||||
total_granted: float | None = None # 总授予额度
|
||||
total_used: float | None = None # 已使用额度
|
||||
total_available: float | None = None # 可用余额
|
||||
expires_at: datetime | None = None # 过期时间
|
||||
currency: str = "USD" # 货币单位
|
||||
extra: Dict[str, Any] = field(default_factory=dict) # 额外信息
|
||||
extra: dict[str, Any] = field(default_factory=dict) # 额外信息
|
||||
|
||||
|
||||
@dataclass
|
||||
class CheckinInfo:
|
||||
"""签到信息"""
|
||||
|
||||
reward: Optional[float] = None # 奖励额度
|
||||
streak_days: Optional[int] = None # 连续签到天数
|
||||
next_reward: Optional[float] = None # 下次奖励
|
||||
message: Optional[str] = None # 签到消息
|
||||
extra: Dict[str, Any] = field(default_factory=dict)
|
||||
reward: float | None = None # 奖励额度
|
||||
streak_days: int | None = None # 连续签到天数
|
||||
next_reward: float | None = None # 下次奖励
|
||||
message: str | None = None # 签到消息
|
||||
extra: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -85,13 +87,13 @@ class ActionResult:
|
||||
|
||||
status: ActionStatus
|
||||
action_type: ProviderActionType
|
||||
data: Optional[Any] = None # 操作返回的数据(如 BalanceInfo, CheckinInfo)
|
||||
message: Optional[str] = None # 消息
|
||||
data: Any | None = None # 操作返回的数据(如 BalanceInfo, CheckinInfo)
|
||||
message: str | None = None # 消息
|
||||
executed_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
|
||||
response_time_ms: Optional[int] = None # 响应时间(毫秒)
|
||||
raw_response: Optional[Dict[str, Any]] = None # 原始响应(调试用)
|
||||
response_time_ms: int | None = None # 响应时间(毫秒)
|
||||
raw_response: dict[str, Any] | None = None # 原始响应(调试用)
|
||||
cache_ttl_seconds: int = 300 # 建议缓存时间
|
||||
retry_after_seconds: Optional[int] = None # 失败后重试间隔
|
||||
retry_after_seconds: int | None = None # 失败后重试间隔
|
||||
|
||||
@property
|
||||
def is_success(self) -> bool:
|
||||
@@ -104,10 +106,10 @@ class ConnectorState:
|
||||
|
||||
status: ConnectorStatus
|
||||
auth_type: ConnectorAuthType
|
||||
connected_at: Optional[datetime] = None
|
||||
expires_at: Optional[datetime] = None
|
||||
last_error: Optional[str] = None
|
||||
extra: Dict[str, Any] = field(default_factory=dict)
|
||||
connected_at: datetime | None = None
|
||||
expires_at: datetime | None = None
|
||||
last_error: str | None = None
|
||||
extra: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -115,21 +117,21 @@ class ProviderOpsConfig:
|
||||
"""Provider 操作配置(存储在 Provider.config['provider_ops'] 中)"""
|
||||
|
||||
architecture_id: str = "generic_api"
|
||||
base_url: Optional[str] = None # API 基础地址
|
||||
base_url: str | None = None # API 基础地址
|
||||
|
||||
# 连接器配置
|
||||
connector_auth_type: ConnectorAuthType = ConnectorAuthType.API_KEY
|
||||
connector_config: Dict[str, Any] = field(default_factory=dict)
|
||||
connector_credentials: Dict[str, Any] = field(default_factory=dict) # 加密存储
|
||||
connector_config: dict[str, Any] = field(default_factory=dict)
|
||||
connector_credentials: dict[str, Any] = field(default_factory=dict) # 加密存储
|
||||
|
||||
# 操作配置
|
||||
actions: Dict[str, Dict[str, Any]] = field(default_factory=dict)
|
||||
actions: dict[str, dict[str, Any]] = field(default_factory=dict)
|
||||
|
||||
# 定时任务配置
|
||||
schedule: Dict[str, str] = field(default_factory=dict) # {action_type: cron_expression}
|
||||
schedule: dict[str, str] = field(default_factory=dict) # {action_type: cron_expression}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Optional[Dict[str, Any]]) -> "ProviderOpsConfig":
|
||||
def from_dict(cls, data: dict[str, Any] | None) -> ProviderOpsConfig:
|
||||
"""从字典创建配置"""
|
||||
if not data:
|
||||
return cls()
|
||||
@@ -146,7 +148,7 @@ class ProviderOpsConfig:
|
||||
schedule=data.get("schedule", {}),
|
||||
)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""转换为字典(用于存储)"""
|
||||
return {
|
||||
"architecture_id": self.architecture_id,
|
||||
|
||||
@@ -13,9 +13,8 @@
|
||||
import statistics
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
from typing import TYPE_CHECKING, Any, Dict, Optional
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from src.core.logger import logger
|
||||
|
||||
from src.config.constants import AdaptiveReservationDefaults
|
||||
|
||||
@@ -68,7 +67,7 @@ class ReservationResult:
|
||||
phase: str # 当前阶段: "probe" | "stable"
|
||||
confidence: float # 置信度 (0-1)
|
||||
load_factor: float # 负载因子 (0-1)
|
||||
details: Dict[str, Any] # 详细信息
|
||||
details: dict[str, Any] # 详细信息
|
||||
|
||||
|
||||
class AdaptiveReservationManager:
|
||||
@@ -91,15 +90,15 @@ class AdaptiveReservationManager:
|
||||
- 调整历史稳定性:最近调整的方差越小越稳定
|
||||
"""
|
||||
|
||||
def __init__(self, config: Optional[ReservationConfig] = None):
|
||||
def __init__(self, config: ReservationConfig | None = None):
|
||||
self.config = config or ReservationConfig()
|
||||
self._cache: Dict[str, ReservationResult] = {} # 简单的内存缓存
|
||||
self._cache: dict[str, ReservationResult] = {} # 简单的内存缓存
|
||||
|
||||
def calculate_reservation(
|
||||
self,
|
||||
key: "ProviderAPIKey",
|
||||
key: ProviderAPIKey,
|
||||
current_usage: int = 0,
|
||||
effective_limit: Optional[int] = None,
|
||||
effective_limit: int | None = None,
|
||||
) -> ReservationResult:
|
||||
"""
|
||||
计算当前应使用的预留比例
|
||||
@@ -148,7 +147,7 @@ class AdaptiveReservationManager:
|
||||
},
|
||||
)
|
||||
|
||||
def _get_total_requests(self, key: "ProviderAPIKey") -> int:
|
||||
def _get_total_requests(self, key: ProviderAPIKey) -> int:
|
||||
"""获取总请求数(用于判断是否过了探测阶段)"""
|
||||
# 使用总请求计数作为基准
|
||||
request_count = key.request_count or 0
|
||||
@@ -165,14 +164,14 @@ class AdaptiveReservationManager:
|
||||
return request_count
|
||||
|
||||
def _calculate_load_ratio(
|
||||
self, current_usage: int, effective_limit: Optional[int]
|
||||
self, current_usage: int, effective_limit: int | None
|
||||
) -> float:
|
||||
"""计算当前负载率"""
|
||||
if not effective_limit or effective_limit <= 0:
|
||||
return 0.0
|
||||
return min(current_usage / effective_limit, 1.0)
|
||||
|
||||
def _calculate_confidence(self, key: "ProviderAPIKey") -> float:
|
||||
def _calculate_confidence(self, key: ProviderAPIKey) -> float:
|
||||
"""
|
||||
计算学习值的置信度 (0-1)
|
||||
|
||||
@@ -186,7 +185,7 @@ class AdaptiveReservationManager:
|
||||
scores["success_score"] + scores["cooldown_score"] + scores["stability_score"], 1.0
|
||||
)
|
||||
|
||||
def _get_confidence_breakdown(self, key: "ProviderAPIKey") -> Dict[str, float]:
|
||||
def _get_confidence_breakdown(self, key: ProviderAPIKey) -> dict[str, float]:
|
||||
"""获取置信度各因素的详细分数"""
|
||||
# 因素1: 成功率(权重 40%)
|
||||
# 使用成功率而非连续成功次数,更准确反映 Key 的稳定性
|
||||
@@ -308,7 +307,7 @@ class AdaptiveReservationManager:
|
||||
|
||||
return f"置信度{confidence:.0%},负载{load_ratio:.0%},动态计算预留"
|
||||
|
||||
def get_stats(self) -> Dict[str, Any]:
|
||||
def get_stats(self) -> dict[str, Any]:
|
||||
"""获取管理器统计信息"""
|
||||
return {
|
||||
"config": {
|
||||
@@ -323,7 +322,7 @@ class AdaptiveReservationManager:
|
||||
|
||||
|
||||
# 全局单例
|
||||
_reservation_manager: Optional[AdaptiveReservationManager] = None
|
||||
_reservation_manager: AdaptiveReservationManager | None = None
|
||||
|
||||
|
||||
def get_adaptive_reservation_manager() -> AdaptiveReservationManager:
|
||||
|
||||
@@ -14,7 +14,7 @@
|
||||
"""
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Dict, List, Optional, cast
|
||||
from typing import Any, cast
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
@@ -90,7 +90,7 @@ class AdaptiveRPMManager:
|
||||
db: Session,
|
||||
key: ProviderAPIKey,
|
||||
rate_limit_info: RateLimitInfo,
|
||||
current_rpm: Optional[int] = None,
|
||||
current_rpm: int | None = None,
|
||||
) -> int:
|
||||
"""
|
||||
处理429错误,调整 RPM 限制
|
||||
@@ -193,7 +193,7 @@ class AdaptiveRPMManager:
|
||||
db: Session,
|
||||
key: ProviderAPIKey,
|
||||
current_rpm: int,
|
||||
) -> Optional[int]:
|
||||
) -> int | None:
|
||||
"""
|
||||
处理成功请求,基于滑动窗口利用率考虑增加 RPM 限制
|
||||
|
||||
@@ -293,7 +293,7 @@ class AdaptiveRPMManager:
|
||||
|
||||
def _update_utilization_window(
|
||||
self, key: ProviderAPIKey, now_ts: float, utilization: float
|
||||
) -> List[Dict[str, Any]]:
|
||||
) -> list[dict[str, Any]]:
|
||||
"""
|
||||
更新利用率滑动窗口
|
||||
|
||||
@@ -305,7 +305,7 @@ class AdaptiveRPMManager:
|
||||
Returns:
|
||||
更新后的采样列表
|
||||
"""
|
||||
samples: List[Dict[str, Any]] = list(key.utilization_samples or [])
|
||||
samples: list[dict[str, Any]] = list(key.utilization_samples or [])
|
||||
|
||||
# 添加新采样
|
||||
samples.append({"ts": now_ts, "util": round(utilization, 3)})
|
||||
@@ -326,10 +326,10 @@ class AdaptiveRPMManager:
|
||||
def _check_increase_conditions(
|
||||
self,
|
||||
key: ProviderAPIKey,
|
||||
samples: List[Dict[str, Any]],
|
||||
samples: list[dict[str, Any]],
|
||||
now: datetime,
|
||||
known_boundary: Optional[int] = None,
|
||||
) -> Optional[str]:
|
||||
known_boundary: int | None = None,
|
||||
) -> str | None:
|
||||
"""
|
||||
检查是否满足扩容条件
|
||||
|
||||
@@ -371,7 +371,7 @@ class AdaptiveRPMManager:
|
||||
return None
|
||||
|
||||
def _should_probe_increase(
|
||||
self, key: ProviderAPIKey, samples: List[Dict[str, Any]], now: datetime
|
||||
self, key: ProviderAPIKey, samples: list[dict[str, Any]], now: datetime
|
||||
) -> bool:
|
||||
"""
|
||||
检查是否应该进行探测性扩容
|
||||
@@ -439,7 +439,7 @@ class AdaptiveRPMManager:
|
||||
def _decrease_limit(
|
||||
self,
|
||||
current_limit: int,
|
||||
current_rpm: Optional[int] = None,
|
||||
current_rpm: int | None = None,
|
||||
) -> int:
|
||||
"""
|
||||
减少 RPM 限制(基于边界记忆策略)
|
||||
@@ -469,7 +469,7 @@ class AdaptiveRPMManager:
|
||||
def _increase_limit(
|
||||
self,
|
||||
current_limit: int,
|
||||
known_boundary: Optional[int] = None,
|
||||
known_boundary: int | None = None,
|
||||
is_probe: bool = False,
|
||||
) -> int:
|
||||
"""
|
||||
@@ -523,7 +523,7 @@ class AdaptiveRPMManager:
|
||||
reason: 调整原因
|
||||
**extra_data: 额外数据
|
||||
"""
|
||||
history: List[Dict[str, Any]] = list(key.adjustment_history or [])
|
||||
history: list[dict[str, Any]] = list(key.adjustment_history or [])
|
||||
|
||||
record = {
|
||||
"timestamp": datetime.now(timezone.utc).isoformat(),
|
||||
@@ -540,7 +540,7 @@ class AdaptiveRPMManager:
|
||||
|
||||
key.adjustment_history = history # type: ignore[assignment]
|
||||
|
||||
def get_adjustment_stats(self, key: ProviderAPIKey) -> Dict[str, Any]:
|
||||
def get_adjustment_stats(self, key: ProviderAPIKey) -> dict[str, Any]:
|
||||
"""
|
||||
获取调整统计信息
|
||||
|
||||
@@ -550,8 +550,8 @@ class AdaptiveRPMManager:
|
||||
Returns:
|
||||
统计信息
|
||||
"""
|
||||
history: List[Dict[str, Any]] = list(key.adjustment_history or [])
|
||||
samples: List[Dict[str, Any]] = list(key.utilization_samples or [])
|
||||
history: list[dict[str, Any]] = list(key.adjustment_history or [])
|
||||
samples: list[dict[str, Any]] = list(key.utilization_samples or [])
|
||||
|
||||
# rpm_limit=NULL 表示自适应,否则为固定限制
|
||||
is_adaptive = key.rpm_limit is None
|
||||
@@ -559,18 +559,18 @@ class AdaptiveRPMManager:
|
||||
effective_limit = current_limit if is_adaptive else int(key.rpm_limit) # type: ignore
|
||||
|
||||
# 计算窗口统计
|
||||
avg_utilization: Optional[float] = None
|
||||
high_util_ratio: Optional[float] = None
|
||||
avg_utilization: float | None = None
|
||||
high_util_ratio: float | None = None
|
||||
if samples:
|
||||
avg_utilization = sum(s["util"] for s in samples) / len(samples)
|
||||
high_util_count = sum(1 for s in samples if s["util"] >= self.UTILIZATION_THRESHOLD)
|
||||
high_util_ratio = high_util_count / len(samples)
|
||||
|
||||
last_429_at_str: Optional[str] = None
|
||||
last_429_at_str: str | None = None
|
||||
if key.last_429_at:
|
||||
last_429_at_str = cast(datetime, key.last_429_at).isoformat()
|
||||
|
||||
last_probe_at_str: Optional[str] = None
|
||||
last_probe_at_str: str | None = None
|
||||
if key.last_probe_increase_at:
|
||||
last_probe_at_str = cast(datetime, key.last_probe_increase_at).isoformat()
|
||||
|
||||
@@ -627,7 +627,7 @@ class AdaptiveRPMManager:
|
||||
|
||||
|
||||
# 全局单例
|
||||
_adaptive_rpm_manager: Optional[AdaptiveRPMManager] = None
|
||||
_adaptive_rpm_manager: AdaptiveRPMManager | None = None
|
||||
|
||||
|
||||
def get_adaptive_rpm_manager() -> AdaptiveRPMManager:
|
||||
|
||||
@@ -8,12 +8,13 @@ RPM 限制管理器 - 支持 Redis 或内存的 Key 级别 RPM 限制
|
||||
4. 支持缓存用户优先级(预留槽位机制)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import math
|
||||
import os
|
||||
import time
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Optional
|
||||
|
||||
import redis.asyncio as aioredis
|
||||
|
||||
@@ -24,8 +25,8 @@ from src.core.logger import logger
|
||||
class ConcurrencyManager:
|
||||
"""Key RPM 限制管理器"""
|
||||
|
||||
_instance: Optional["ConcurrencyManager"] = None
|
||||
_redis: Optional[aioredis.Redis] = None
|
||||
_instance: ConcurrencyManager | None = None
|
||||
_redis: aioredis.Redis | None = None
|
||||
_key_rpm_bucket_seconds: int = 60
|
||||
_key_rpm_key_ttl_seconds: int = 120 # 2 分钟,足够覆盖当前分钟与边界
|
||||
|
||||
@@ -47,7 +48,7 @@ class ConcurrencyManager:
|
||||
self._last_cleanup_bucket: int = 0 # 上次清理时的 bucket,用于定期清理过期数据
|
||||
self._last_cleanup_time: float = 0 # 上次清理的时间戳,用于强制定期清理
|
||||
self._cleanup_interval_seconds: int = 300 # 强制清理间隔(5 分钟)
|
||||
self._cleanup_task: Optional[asyncio.Task] = None # 后台清理任务
|
||||
self._cleanup_task: asyncio.Task | None = None # 后台清理任务
|
||||
|
||||
# 内存模式下的最大条目限制,防止内存泄漏(支持环境变量覆盖)
|
||||
self._max_memory_rpm_entries: int = int(
|
||||
@@ -127,13 +128,13 @@ class ConcurrencyManager:
|
||||
self._owns_redis = False
|
||||
|
||||
@classmethod
|
||||
def _get_rpm_bucket(cls, now_ts: Optional[float] = None) -> int:
|
||||
def _get_rpm_bucket(cls, now_ts: float | None = None) -> int:
|
||||
"""获取当前 RPM 计数桶(按分钟)"""
|
||||
ts = now_ts if now_ts is not None else time.time()
|
||||
return int(ts // cls._key_rpm_bucket_seconds)
|
||||
|
||||
@classmethod
|
||||
def _get_key_key(cls, key_id: str, bucket: Optional[int] = None) -> str:
|
||||
def _get_key_key(cls, key_id: str, bucket: int | None = None) -> str:
|
||||
"""获取 ProviderAPIKey RPM 计数的 Redis Key(按分钟桶)"""
|
||||
b = bucket if bucket is not None else cls._get_rpm_bucket()
|
||||
return f"rpm:key:{key_id}:{b}"
|
||||
@@ -261,9 +262,9 @@ class ConcurrencyManager:
|
||||
async def check_rpm_available(
|
||||
self,
|
||||
key_id: str,
|
||||
key_rpm_limit: Optional[int],
|
||||
key_rpm_limit: int | None,
|
||||
is_cached_user: bool = False,
|
||||
cache_reservation_ratio: Optional[float] = None,
|
||||
cache_reservation_ratio: float | None = None,
|
||||
) -> bool:
|
||||
"""
|
||||
检查是否可以通过 RPM 限制(不实际增加计数)
|
||||
@@ -298,9 +299,9 @@ class ConcurrencyManager:
|
||||
async def acquire_rpm_slot(
|
||||
self,
|
||||
key_id: str,
|
||||
key_rpm_limit: Optional[int],
|
||||
key_rpm_limit: int | None,
|
||||
is_cached_user: bool = False,
|
||||
cache_reservation_ratio: Optional[float] = None,
|
||||
cache_reservation_ratio: float | None = None,
|
||||
) -> bool:
|
||||
"""
|
||||
尝试获取 RPM 槽位(支持缓存用户优先级)
|
||||
@@ -448,9 +449,9 @@ class ConcurrencyManager:
|
||||
async def rpm_guard(
|
||||
self,
|
||||
key_id: str,
|
||||
key_rpm_limit: Optional[int],
|
||||
key_rpm_limit: int | None,
|
||||
is_cached_user: bool = False,
|
||||
cache_reservation_ratio: Optional[float] = None,
|
||||
cache_reservation_ratio: float | None = None,
|
||||
):
|
||||
"""
|
||||
RPM 限制上下文管理器(支持缓存用户优先级)
|
||||
@@ -595,7 +596,7 @@ class ConcurrencyManager:
|
||||
|
||||
|
||||
# 全局单例
|
||||
_concurrency_manager: Optional[ConcurrencyManager] = None
|
||||
_concurrency_manager: ConcurrencyManager | None = None
|
||||
|
||||
|
||||
async def get_concurrency_manager() -> ConcurrencyManager:
|
||||
|
||||
@@ -3,7 +3,6 @@
|
||||
"""
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from typing import Dict, Optional
|
||||
|
||||
from src.core.logger import logger
|
||||
|
||||
@@ -24,12 +23,12 @@ class RateLimitInfo:
|
||||
def __init__(
|
||||
self,
|
||||
limit_type: str,
|
||||
retry_after: Optional[int] = None,
|
||||
limit_value: Optional[int] = None,
|
||||
remaining: Optional[int] = None,
|
||||
reset_at: Optional[datetime] = None,
|
||||
current_usage: Optional[int] = None,
|
||||
raw_headers: Optional[Dict[str, str]] = None,
|
||||
retry_after: int | None = None,
|
||||
limit_value: int | None = None,
|
||||
remaining: int | None = None,
|
||||
reset_at: datetime | None = None,
|
||||
current_usage: int | None = None,
|
||||
raw_headers: dict[str, str] | None = None,
|
||||
):
|
||||
self.limit_type = limit_type
|
||||
self.retry_after = retry_after # 需要等待的秒数
|
||||
@@ -60,9 +59,9 @@ class RateLimitDetector:
|
||||
|
||||
@staticmethod
|
||||
def detect_from_headers(
|
||||
headers: Dict[str, str],
|
||||
headers: dict[str, str],
|
||||
provider_name: str = "unknown",
|
||||
current_usage: Optional[int] = None,
|
||||
current_usage: int | None = None,
|
||||
) -> RateLimitInfo:
|
||||
"""
|
||||
从响应头中检测速率限制类型
|
||||
@@ -88,8 +87,8 @@ class RateLimitDetector:
|
||||
|
||||
@staticmethod
|
||||
def _parse_anthropic_headers(
|
||||
headers: Dict[str, str],
|
||||
current_usage: Optional[int] = None,
|
||||
headers: dict[str, str],
|
||||
current_usage: int | None = None,
|
||||
) -> RateLimitInfo:
|
||||
"""
|
||||
解析 Anthropic Claude API 的速率限制头
|
||||
@@ -195,8 +194,8 @@ class RateLimitDetector:
|
||||
|
||||
@staticmethod
|
||||
def _parse_openai_headers(
|
||||
headers: Dict[str, str],
|
||||
current_usage: Optional[int] = None,
|
||||
headers: dict[str, str],
|
||||
current_usage: int | None = None,
|
||||
) -> RateLimitInfo:
|
||||
"""
|
||||
解析 OpenAI API 的速率限制头
|
||||
@@ -289,8 +288,8 @@ class RateLimitDetector:
|
||||
|
||||
@staticmethod
|
||||
def _parse_generic_headers(
|
||||
headers: Dict[str, str],
|
||||
current_usage: Optional[int] = None,
|
||||
headers: dict[str, str],
|
||||
current_usage: int | None = None,
|
||||
) -> RateLimitInfo:
|
||||
"""
|
||||
解析通用的速率限制头
|
||||
@@ -371,7 +370,7 @@ class RateLimitDetector:
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _parse_retry_after(headers: Dict[str, str]) -> Optional[int]:
|
||||
def _parse_retry_after(headers: dict[str, str]) -> int | None:
|
||||
"""解析 Retry-After 头"""
|
||||
retry_after_str = headers.get("retry-after")
|
||||
if not retry_after_str:
|
||||
@@ -390,7 +389,7 @@ class RateLimitDetector:
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _parse_int(value: Optional[str]) -> Optional[int]:
|
||||
def _parse_int(value: str | None) -> int | None:
|
||||
"""安全解析整数"""
|
||||
if not value:
|
||||
return None
|
||||
@@ -400,7 +399,7 @@ class RateLimitDetector:
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _parse_datetime(value: Optional[str]) -> Optional[datetime]:
|
||||
def _parse_datetime(value: str | None) -> datetime | None:
|
||||
"""安全解析ISO 8601日期时间"""
|
||||
if not value:
|
||||
return None
|
||||
@@ -415,9 +414,9 @@ class RateLimitDetector:
|
||||
|
||||
# 便捷函数
|
||||
def detect_rate_limit_type(
|
||||
headers: Dict[str, str],
|
||||
headers: dict[str, str],
|
||||
provider_name: str = "unknown",
|
||||
current_usage: Optional[int] = None,
|
||||
current_usage: int | None = None,
|
||||
) -> RateLimitInfo:
|
||||
"""
|
||||
检测速率限制类型(便捷函数)
|
||||
|
||||
@@ -5,8 +5,6 @@ IP 级别的速率限制服务
|
||||
"""
|
||||
|
||||
import ipaddress
|
||||
from datetime import datetime, timezone
|
||||
from typing import Dict, Optional, Set
|
||||
|
||||
from src.clients.redis_client import get_redis_client
|
||||
from src.core.logger import logger
|
||||
@@ -34,7 +32,7 @@ class IPRateLimiter:
|
||||
|
||||
@staticmethod
|
||||
async def check_limit(
|
||||
ip_address: str, endpoint_type: str = "default", limit: Optional[int] = None
|
||||
ip_address: str, endpoint_type: str = "default", limit: int | None = None
|
||||
) -> tuple[bool, int, int]:
|
||||
"""
|
||||
检查 IP 是否超过速率限制
|
||||
@@ -102,7 +100,7 @@ class IPRateLimiter:
|
||||
|
||||
@staticmethod
|
||||
async def add_to_blacklist(
|
||||
ip_address: str, reason: str = "manual", ttl: Optional[int] = None
|
||||
ip_address: str, reason: str = "manual", ttl: int | None = None
|
||||
) -> bool:
|
||||
"""
|
||||
将 IP 加入黑名单
|
||||
@@ -301,7 +299,7 @@ class IPRateLimiter:
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
async def get_blacklist_stats() -> Dict:
|
||||
async def get_blacklist_stats() -> dict:
|
||||
"""
|
||||
获取黑名单统计信息
|
||||
|
||||
@@ -332,7 +330,7 @@ class IPRateLimiter:
|
||||
return {"available": False, "total": 0, "error": str(e)}
|
||||
|
||||
@staticmethod
|
||||
async def get_whitelist() -> Set[str]:
|
||||
async def get_whitelist() -> set[str]:
|
||||
"""
|
||||
获取白名单列表
|
||||
|
||||
|
||||
@@ -4,7 +4,6 @@
|
||||
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from typing import List, Optional
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
@@ -22,16 +21,16 @@ class RequestCandidateService:
|
||||
request_id: str,
|
||||
candidate_index: int,
|
||||
retry_index: int = 0, # 新增:重试序号
|
||||
user_id: Optional[str] = None,
|
||||
api_key_id: Optional[str] = None,
|
||||
provider_id: Optional[str] = None,
|
||||
endpoint_id: Optional[str] = None,
|
||||
key_id: Optional[str] = None,
|
||||
user_id: str | None = None,
|
||||
api_key_id: str | None = None,
|
||||
provider_id: str | None = None,
|
||||
endpoint_id: str | None = None,
|
||||
key_id: str | None = None,
|
||||
status: str = "available",
|
||||
skip_reason: Optional[str] = None,
|
||||
skip_reason: str | None = None,
|
||||
is_cached: bool = False,
|
||||
extra_data: Optional[dict] = None,
|
||||
required_capabilities: Optional[dict] = None,
|
||||
extra_data: dict | None = None,
|
||||
required_capabilities: dict | None = None,
|
||||
) -> RequestCandidate:
|
||||
"""
|
||||
创建候选记录
|
||||
@@ -116,7 +115,7 @@ class RequestCandidateService:
|
||||
db: Session,
|
||||
candidate_id: str,
|
||||
status_code: int = 200,
|
||||
concurrent_requests: Optional[int] = None,
|
||||
concurrent_requests: int | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
标记候选为流式传输中
|
||||
@@ -144,8 +143,8 @@ class RequestCandidateService:
|
||||
candidate_id: str,
|
||||
status_code: int,
|
||||
latency_ms: int,
|
||||
concurrent_requests: Optional[int] = None,
|
||||
extra_data: Optional[dict] = None,
|
||||
concurrent_requests: int | None = None,
|
||||
extra_data: dict | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
标记候选执行成功
|
||||
@@ -180,10 +179,10 @@ class RequestCandidateService:
|
||||
candidate_id: str,
|
||||
error_type: str,
|
||||
error_message: str,
|
||||
status_code: Optional[int] = None,
|
||||
latency_ms: Optional[int] = None,
|
||||
concurrent_requests: Optional[int] = None,
|
||||
extra_data: Optional[dict] = None,
|
||||
status_code: int | None = None,
|
||||
latency_ms: int | None = None,
|
||||
concurrent_requests: int | None = None,
|
||||
extra_data: dict | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
标记候选执行失败
|
||||
@@ -218,9 +217,9 @@ class RequestCandidateService:
|
||||
db: Session,
|
||||
candidate_id: str,
|
||||
status_code: int = 499,
|
||||
latency_ms: Optional[int] = None,
|
||||
concurrent_requests: Optional[int] = None,
|
||||
extra_data: Optional[dict] = None,
|
||||
latency_ms: int | None = None,
|
||||
concurrent_requests: int | None = None,
|
||||
extra_data: dict | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
标记候选被客户端取消
|
||||
@@ -248,7 +247,7 @@ class RequestCandidateService:
|
||||
|
||||
@staticmethod
|
||||
def mark_candidate_skipped(
|
||||
db: Session, candidate_id: str, skip_reason: Optional[str] = None
|
||||
db: Session, candidate_id: str, skip_reason: str | None = None
|
||||
) -> None:
|
||||
"""
|
||||
标记候选为已跳过
|
||||
@@ -267,7 +266,7 @@ class RequestCandidateService:
|
||||
get_batch_committer().mark_dirty(db)
|
||||
|
||||
@staticmethod
|
||||
def get_candidates_by_request_id(db: Session, request_id: str) -> List[RequestCandidate]:
|
||||
def get_candidates_by_request_id(db: Session, request_id: str) -> list[RequestCandidate]:
|
||||
"""
|
||||
获取请求的所有候选记录
|
||||
|
||||
|
||||
@@ -4,7 +4,9 @@
|
||||
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Callable, Optional, Union
|
||||
from typing import Any
|
||||
|
||||
from collections.abc import Callable
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
@@ -24,12 +26,12 @@ class ExecutionContext:
|
||||
provider_id: str
|
||||
endpoint_id: str
|
||||
key_id: str
|
||||
user_id: Optional[str]
|
||||
api_key_id: Optional[str]
|
||||
user_id: str | None
|
||||
api_key_id: str | None
|
||||
is_cached_user: bool
|
||||
start_time: Optional[float] = None
|
||||
elapsed_ms: Optional[int] = None
|
||||
concurrent_requests: Optional[int] = None
|
||||
start_time: float | None = None
|
||||
elapsed_ms: int | None = None
|
||||
concurrent_requests: int | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -59,8 +61,8 @@ class RequestExecutor:
|
||||
candidate_index: int,
|
||||
user_api_key,
|
||||
request_func: Callable,
|
||||
request_id: Optional[str],
|
||||
api_format: Union[str, APIFormat],
|
||||
request_id: str | None,
|
||||
api_format: str | APIFormat,
|
||||
model_name: str,
|
||||
is_stream: bool = False,
|
||||
) -> ExecutionResult:
|
||||
|
||||
@@ -13,9 +13,11 @@
|
||||
- ChatAdapterBase 使用 RequestResult 处理异常响应
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from typing import Any, AsyncIterator, Dict, Optional
|
||||
from typing import Any
|
||||
from collections.abc import AsyncIterator
|
||||
|
||||
|
||||
class RequestStatus(Enum):
|
||||
@@ -50,20 +52,20 @@ class RequestMetadata:
|
||||
model: str = "unknown"
|
||||
|
||||
# Provider 追踪信息
|
||||
provider_id: Optional[str] = None
|
||||
provider_endpoint_id: Optional[str] = None
|
||||
provider_api_key_id: Optional[str] = None
|
||||
provider_id: str | None = None
|
||||
provider_endpoint_id: str | None = None
|
||||
provider_api_key_id: str | None = None
|
||||
|
||||
# 请求/响应头
|
||||
provider_request_headers: Dict[str, str] = field(default_factory=dict)
|
||||
provider_response_headers: Dict[str, str] = field(default_factory=dict)
|
||||
provider_request_headers: dict[str, str] = field(default_factory=dict)
|
||||
provider_response_headers: dict[str, str] = field(default_factory=dict)
|
||||
|
||||
# 其他元数据
|
||||
attempt_id: Optional[str] = None
|
||||
original_model: Optional[str] = None # 用户请求的原始模型名(用于价格计算)
|
||||
attempt_id: str | None = None
|
||||
original_model: str | None = None # 用户请求的原始模型名(用于价格计算)
|
||||
|
||||
# Provider 响应元数据(存储 provider 返回的额外信息,如 Gemini 的 modelVersion)
|
||||
response_metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
response_metadata: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
def with_provider_info(
|
||||
self,
|
||||
@@ -71,7 +73,7 @@ class RequestMetadata:
|
||||
provider_id: str,
|
||||
provider_endpoint_id: str,
|
||||
provider_api_key_id: str,
|
||||
) -> "RequestMetadata":
|
||||
) -> RequestMetadata:
|
||||
"""返回包含 Provider 信息的新 RequestMetadata"""
|
||||
return RequestMetadata(
|
||||
api_format=self.api_format,
|
||||
@@ -87,7 +89,7 @@ class RequestMetadata:
|
||||
response_metadata=self.response_metadata,
|
||||
)
|
||||
|
||||
def with_response_headers(self, headers: Dict[str, str]) -> "RequestMetadata":
|
||||
def with_response_headers(self, headers: dict[str, str]) -> RequestMetadata:
|
||||
"""返回包含响应头的新 RequestMetadata"""
|
||||
return RequestMetadata(
|
||||
api_format=self.api_format,
|
||||
@@ -151,8 +153,8 @@ class RequestResult:
|
||||
metadata: RequestMetadata
|
||||
|
||||
# 响应相关
|
||||
response_data: Optional[Any] = None # 成功时的响应数据
|
||||
stream: Optional[AsyncIterator[str]] = None # 流式响应
|
||||
response_data: Any | None = None # 成功时的响应数据
|
||||
stream: AsyncIterator[str] | None = None # 流式响应
|
||||
|
||||
# 使用量和费用
|
||||
usage: UsageInfo = field(default_factory=UsageInfo)
|
||||
@@ -160,16 +162,16 @@ class RequestResult:
|
||||
|
||||
# 错误信息
|
||||
status_code: int = 200
|
||||
error_message: Optional[str] = None
|
||||
error_type: Optional[str] = None
|
||||
error_message: str | None = None
|
||||
error_type: str | None = None
|
||||
|
||||
# 计时
|
||||
response_time_ms: int = 0
|
||||
|
||||
# 请求信息(用于记录)
|
||||
is_stream: bool = False
|
||||
request_headers: Dict[str, str] = field(default_factory=dict)
|
||||
request_body: Dict[str, Any] = field(default_factory=dict)
|
||||
request_headers: dict[str, str] = field(default_factory=dict)
|
||||
request_body: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
@property
|
||||
def is_success(self) -> bool:
|
||||
@@ -191,7 +193,7 @@ class RequestResult:
|
||||
usage: UsageInfo,
|
||||
response_time_ms: int,
|
||||
is_stream: bool = False,
|
||||
) -> "RequestResult":
|
||||
) -> RequestResult:
|
||||
"""创建成功的请求结果"""
|
||||
return cls(
|
||||
status=RequestStatus.SUCCESS,
|
||||
@@ -212,7 +214,7 @@ class RequestResult:
|
||||
error_type: str,
|
||||
response_time_ms: int,
|
||||
is_stream: bool = False,
|
||||
) -> "RequestResult":
|
||||
) -> RequestResult:
|
||||
"""创建失败的请求结果"""
|
||||
return cls(
|
||||
status=RequestStatus.FAILED,
|
||||
@@ -229,9 +231,9 @@ class RequestResult:
|
||||
cls,
|
||||
metadata: RequestMetadata,
|
||||
response_time_ms: int,
|
||||
usage: Optional[UsageInfo] = None,
|
||||
usage: UsageInfo | None = None,
|
||||
is_stream: bool = False,
|
||||
) -> "RequestResult":
|
||||
) -> RequestResult:
|
||||
"""创建客户端取消的请求结果"""
|
||||
return cls(
|
||||
status=RequestStatus.CANCELLED,
|
||||
@@ -252,7 +254,7 @@ class RequestResult:
|
||||
model: str,
|
||||
response_time_ms: int,
|
||||
is_stream: bool = False,
|
||||
) -> "RequestResult":
|
||||
) -> RequestResult:
|
||||
"""从异常创建失败的请求结果"""
|
||||
# 尝试从异常中提取 metadata
|
||||
existing_metadata = getattr(exception, "request_metadata", None)
|
||||
@@ -338,7 +340,7 @@ class StreamWithMetadata:
|
||||
self,
|
||||
stream: AsyncIterator[str],
|
||||
metadata: RequestMetadata,
|
||||
response_headers_container: Optional[Dict[str, Any]] = None,
|
||||
response_headers_container: dict[str, Any] | None = None,
|
||||
):
|
||||
self.stream = stream
|
||||
self.metadata = metadata
|
||||
|
||||
@@ -3,9 +3,8 @@
|
||||
"""
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from typing import List, Optional
|
||||
|
||||
from sqlalchemy import and_, or_
|
||||
from sqlalchemy import or_
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.exceptions import ForbiddenException, NotFoundException
|
||||
@@ -26,8 +25,8 @@ class AnnouncementService:
|
||||
type: str = "info",
|
||||
priority: int = 0,
|
||||
is_pinned: bool = False,
|
||||
start_time: Optional[datetime] = None,
|
||||
end_time: Optional[datetime] = None,
|
||||
start_time: datetime | None = None,
|
||||
end_time: datetime | None = None,
|
||||
) -> Announcement:
|
||||
"""创建公告"""
|
||||
# 验证作者是否为管理员
|
||||
@@ -61,7 +60,7 @@ class AnnouncementService:
|
||||
@staticmethod
|
||||
def get_announcements(
|
||||
db: Session,
|
||||
user_id: Optional[str] = None, # UUID
|
||||
user_id: str | None = None, # UUID
|
||||
active_only: bool = True,
|
||||
include_read_status: bool = False,
|
||||
limit: int = 50,
|
||||
@@ -148,14 +147,14 @@ class AnnouncementService:
|
||||
db: Session,
|
||||
announcement_id: str, # UUID
|
||||
user_id: str, # UUID
|
||||
title: Optional[str] = None,
|
||||
content: Optional[str] = None,
|
||||
type: Optional[str] = None,
|
||||
priority: Optional[int] = None,
|
||||
is_active: Optional[bool] = None,
|
||||
is_pinned: Optional[bool] = None,
|
||||
start_time: Optional[datetime] = None,
|
||||
end_time: Optional[datetime] = None,
|
||||
title: str | None = None,
|
||||
content: str | None = None,
|
||||
type: str | None = None,
|
||||
priority: int | None = None,
|
||||
is_active: bool | None = None,
|
||||
is_pinned: bool | None = None,
|
||||
start_time: datetime | None = None,
|
||||
end_time: datetime | None = None,
|
||||
) -> Announcement:
|
||||
"""更新公告"""
|
||||
# 验证用户是否为管理员
|
||||
@@ -230,7 +229,7 @@ class AnnouncementService:
|
||||
logger.info(f"User {user_id} marked announcement {announcement_id} as read")
|
||||
|
||||
@staticmethod
|
||||
def get_active_announcements(db: Session, user_id: Optional[str] = None) -> dict: # UUID
|
||||
def get_active_announcements(db: Session, user_id: str | None = None) -> dict: # UUID
|
||||
"""获取当前有效的公告(首页展示用)"""
|
||||
return AnnouncementService.get_announcements(
|
||||
db=db,
|
||||
|
||||
@@ -4,7 +4,7 @@
|
||||
"""
|
||||
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Any, Dict, List, Optional
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
@@ -29,14 +29,14 @@ class AuditService:
|
||||
db: Session,
|
||||
event_type: AuditEventType,
|
||||
description: str,
|
||||
user_id: Optional[str] = None, # UUID
|
||||
api_key_id: Optional[str] = None, # UUID
|
||||
ip_address: Optional[str] = None,
|
||||
user_agent: Optional[str] = None,
|
||||
request_id: Optional[str] = None,
|
||||
status_code: Optional[int] = None,
|
||||
error_message: Optional[str] = None,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
user_id: str | None = None, # UUID
|
||||
api_key_id: str | None = None, # UUID
|
||||
ip_address: str | None = None,
|
||||
user_agent: str | None = None,
|
||||
request_id: str | None = None,
|
||||
status_code: int | None = None,
|
||||
error_message: str | None = None,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
) -> AuditLog:
|
||||
"""
|
||||
记录审计事件
|
||||
@@ -106,8 +106,8 @@ class AuditService:
|
||||
success: bool,
|
||||
ip_address: str,
|
||||
user_agent: str,
|
||||
user_id: Optional[str] = None, # UUID
|
||||
error_reason: Optional[str] = None,
|
||||
user_id: str | None = None, # UUID
|
||||
error_reason: str | None = None,
|
||||
):
|
||||
"""
|
||||
记录登录尝试
|
||||
@@ -147,10 +147,10 @@ class AuditService:
|
||||
success: bool,
|
||||
ip_address: str,
|
||||
status_code: int,
|
||||
error_message: Optional[str] = None,
|
||||
input_tokens: Optional[int] = None,
|
||||
output_tokens: Optional[int] = None,
|
||||
cost_usd: Optional[float] = None,
|
||||
error_message: str | None = None,
|
||||
input_tokens: int | None = None,
|
||||
output_tokens: int | None = None,
|
||||
cost_usd: float | None = None,
|
||||
):
|
||||
"""
|
||||
记录API请求
|
||||
@@ -201,9 +201,9 @@ class AuditService:
|
||||
event_type: AuditEventType,
|
||||
description: str,
|
||||
ip_address: str,
|
||||
user_id: Optional[str] = None, # UUID
|
||||
user_id: str | None = None, # UUID
|
||||
severity: str = "medium",
|
||||
details: Optional[Dict[str, Any]] = None,
|
||||
details: dict[str, Any] | None = None,
|
||||
):
|
||||
"""
|
||||
记录安全事件
|
||||
@@ -238,9 +238,9 @@ class AuditService:
|
||||
def get_user_audit_logs(
|
||||
db: Session,
|
||||
user_id: str, # UUID
|
||||
event_types: Optional[List[AuditEventType]] = None,
|
||||
event_types: list[AuditEventType] | None = None,
|
||||
limit: int = 100,
|
||||
) -> List[AuditLog]:
|
||||
) -> list[AuditLog]:
|
||||
"""
|
||||
获取用户的审计日志
|
||||
|
||||
@@ -262,7 +262,7 @@ class AuditService:
|
||||
return query.order_by(AuditLog.created_at.desc()).limit(limit).all()
|
||||
|
||||
@staticmethod
|
||||
def get_suspicious_activities(db: Session, hours: int = 24, limit: int = 100) -> List[AuditLog]:
|
||||
def get_suspicious_activities(db: Session, hours: int = 24, limit: int = 100) -> list[AuditLog]:
|
||||
"""
|
||||
获取可疑活动
|
||||
|
||||
@@ -292,7 +292,7 @@ class AuditService:
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def analyze_user_behavior(db: Session, user_id: str, days: int = 30) -> Dict[str, Any]: # UUID
|
||||
def analyze_user_behavior(db: Session, user_id: str, days: int = 30) -> dict[str, Any]: # UUID
|
||||
"""
|
||||
分析用户行为
|
||||
|
||||
@@ -373,16 +373,16 @@ class AuditService:
|
||||
def log_event_auto(
|
||||
event_type: AuditEventType,
|
||||
description: str,
|
||||
user_id: Optional[str] = None,
|
||||
api_key_id: Optional[str] = None,
|
||||
ip_address: Optional[str] = None,
|
||||
user_agent: Optional[str] = None,
|
||||
request_id: Optional[str] = None,
|
||||
status_code: Optional[int] = None,
|
||||
error_message: Optional[str] = None,
|
||||
event_metadata: Optional[Dict[str, Any]] = None,
|
||||
db: Optional[Session] = None,
|
||||
) -> Optional[AuditLog]:
|
||||
user_id: str | None = None,
|
||||
api_key_id: str | None = None,
|
||||
ip_address: str | None = None,
|
||||
user_agent: str | None = None,
|
||||
request_id: str | None = None,
|
||||
status_code: int | None = None,
|
||||
error_message: str | None = None,
|
||||
event_metadata: dict[str, Any] | None = None,
|
||||
db: Session | None = None,
|
||||
) -> AuditLog | None:
|
||||
"""
|
||||
自动管理数据库会话的审计日志记录方法
|
||||
适用于中间件等无法直接获取数据库会话的场景
|
||||
|
||||
@@ -12,7 +12,7 @@
|
||||
import asyncio
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Dict
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
@@ -26,7 +26,7 @@ class WarmupContext:
|
||||
|
||||
db: Session
|
||||
user: Any # User model
|
||||
audit_metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
audit_metadata: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
def add_audit_metadata(self, **kwargs: Any) -> None:
|
||||
"""兼容 ApiRequestContext 接口"""
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
import json
|
||||
import time
|
||||
from enum import Enum
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
@@ -25,10 +25,10 @@ class LogLevel(str, Enum):
|
||||
_CONFIG_CACHE_TTL = 60 # 1 分钟
|
||||
|
||||
# 进程内缓存存储: {key: (value, expire_time)}
|
||||
_config_cache: Dict[str, Tuple[Any, float]] = {}
|
||||
_config_cache: dict[str, tuple[Any, float]] = {}
|
||||
|
||||
|
||||
def _get_cached_config(key: str) -> Tuple[bool, Any]:
|
||||
def _get_cached_config(key: str) -> tuple[bool, Any]:
|
||||
"""从进程内缓存获取配置值
|
||||
|
||||
Returns:
|
||||
@@ -48,7 +48,7 @@ def _set_cached_config(key: str, value: Any) -> None:
|
||||
_config_cache[key] = (value, time.time() + _CONFIG_CACHE_TTL)
|
||||
|
||||
|
||||
def invalidate_config_cache(key: Optional[str] = None) -> None:
|
||||
def invalidate_config_cache(key: str | None = None) -> None:
|
||||
"""清除配置缓存
|
||||
|
||||
Args:
|
||||
@@ -174,7 +174,7 @@ class SystemConfigService:
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def get_config(cls, db: Session, key: str, default: Any = None) -> Optional[Any]:
|
||||
def get_config(cls, db: Session, key: str, default: Any = None) -> Any | None:
|
||||
"""获取系统配置值(带进程内缓存)"""
|
||||
# 1. 检查进程内缓存
|
||||
hit, cached_value = _get_cached_config(key)
|
||||
@@ -196,7 +196,7 @@ class SystemConfigService:
|
||||
return default
|
||||
|
||||
@classmethod
|
||||
def get_configs(cls, db: Session, keys: List[str]) -> Dict[str, Any]:
|
||||
def get_configs(cls, db: Session, keys: list[str]) -> dict[str, Any]:
|
||||
"""
|
||||
批量获取系统配置值
|
||||
|
||||
@@ -248,7 +248,7 @@ class SystemConfigService:
|
||||
return config
|
||||
|
||||
@staticmethod
|
||||
def get_default_provider(db: Session) -> Optional[str]:
|
||||
def get_default_provider(db: Session) -> str | None:
|
||||
"""
|
||||
获取系统默认提供商
|
||||
优先级:1. 管理员设置的默认提供商 2. 数据库中第一个可用提供商
|
||||
@@ -355,7 +355,7 @@ class SystemConfigService:
|
||||
return cls.get_config(db, "sensitive_headers", [])
|
||||
|
||||
@classmethod
|
||||
def mask_sensitive_headers(cls, db: Session, headers: Dict[str, Any]) -> Dict[str, Any]:
|
||||
def mask_sensitive_headers(cls, db: Session, headers: dict[str, Any]) -> dict[str, Any]:
|
||||
"""脱敏敏感请求头"""
|
||||
if not cls.should_mask_sensitive_data(db):
|
||||
return headers
|
||||
|
||||
@@ -6,9 +6,10 @@
|
||||
数据存储仍然使用 UTC。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from datetime import datetime
|
||||
from typing import Optional
|
||||
|
||||
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
||||
from apscheduler.triggers.cron import CronTrigger
|
||||
@@ -23,14 +24,14 @@ APP_TIMEZONE = os.getenv("APP_TIMEZONE", "Asia/Shanghai")
|
||||
class TaskScheduler:
|
||||
"""统一定时任务调度器"""
|
||||
|
||||
_instance: Optional["TaskScheduler"] = None
|
||||
_instance: TaskScheduler | None = None
|
||||
|
||||
def __init__(self):
|
||||
self.scheduler = AsyncIOScheduler(timezone=APP_TIMEZONE)
|
||||
self._started = False
|
||||
|
||||
@classmethod
|
||||
def get_instance(cls) -> "TaskScheduler":
|
||||
def get_instance(cls) -> TaskScheduler:
|
||||
"""获取调度器单例"""
|
||||
if cls._instance is None:
|
||||
cls._instance = TaskScheduler()
|
||||
|
||||
@@ -6,7 +6,6 @@
|
||||
import os
|
||||
import uuid
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Optional
|
||||
|
||||
from sqlalchemy import and_, func
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
@@ -3,7 +3,6 @@ API密钥统计同步服务
|
||||
定期同步API密钥的统计数据,确保与实际使用记录一致
|
||||
"""
|
||||
|
||||
from typing import Optional
|
||||
|
||||
from sqlalchemy import func
|
||||
from sqlalchemy.orm import Session
|
||||
@@ -20,7 +19,7 @@ class SyncStatsService:
|
||||
BATCH_SIZE = 100
|
||||
|
||||
@staticmethod
|
||||
def sync_api_key_stats(db: Session, api_key_id: Optional[str] = None) -> dict: # UUID
|
||||
def sync_api_key_stats(db: Session, api_key_id: str | None = None) -> dict: # UUID
|
||||
"""
|
||||
同步API密钥的统计数据
|
||||
|
||||
|
||||
@@ -4,14 +4,13 @@ Usage Redis Streams consumer.
|
||||
高性能消费者实现,支持批量处理和单次提交多条记录。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import socket
|
||||
import time
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
from typing import Any
|
||||
|
||||
from redis.exceptions import ResponseError
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
@@ -47,7 +46,7 @@ def _parse_body(value: Any) -> Any:
|
||||
return value
|
||||
|
||||
|
||||
def _event_to_record(event: UsageEvent) -> Dict[str, Any]:
|
||||
def _event_to_record(event: UsageEvent) -> dict[str, Any]:
|
||||
"""将 UsageEvent 转换为 record_usage_batch 所需的字典格式"""
|
||||
data = event.data
|
||||
status = "completed"
|
||||
@@ -121,7 +120,7 @@ class UsageQueueConsumer:
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._running = False
|
||||
self._task: Optional[asyncio.Task] = None
|
||||
self._task: asyncio.Task | None = None
|
||||
self._consumer = _consumer_name()
|
||||
self._last_claim = 0.0
|
||||
self._last_metrics_log = 0.0
|
||||
@@ -220,9 +219,9 @@ class UsageQueueConsumer:
|
||||
return
|
||||
|
||||
# 分类消息
|
||||
streaming_messages: List[Tuple[str, UsageEvent]] = []
|
||||
record_messages: List[Tuple[str, Dict[str, Any], UsageEvent]] = []
|
||||
failed_messages: List[Tuple[str, Dict[str, Any], Exception]] = []
|
||||
streaming_messages: list[tuple[str, UsageEvent]] = []
|
||||
record_messages: list[tuple[str, dict[str, Any], UsageEvent]] = []
|
||||
failed_messages: list[tuple[str, dict[str, Any], Exception]] = []
|
||||
|
||||
for message_id, fields in messages:
|
||||
try:
|
||||
@@ -249,10 +248,10 @@ class UsageQueueConsumer:
|
||||
async def _process_streaming_batch(
|
||||
self,
|
||||
redis_client,
|
||||
messages: List[Tuple[str, UsageEvent]],
|
||||
messages: list[tuple[str, UsageEvent]],
|
||||
) -> None:
|
||||
"""批量处理 STREAMING 事件(状态更新)"""
|
||||
success_ids: List[str] = []
|
||||
success_ids: list[str] = []
|
||||
|
||||
for message_id, event in messages:
|
||||
try:
|
||||
@@ -271,16 +270,16 @@ class UsageQueueConsumer:
|
||||
async def _process_record_batch(
|
||||
self,
|
||||
redis_client,
|
||||
messages: List[Tuple[str, Dict[str, Any], UsageEvent]],
|
||||
messages: list[tuple[str, dict[str, Any], UsageEvent]],
|
||||
) -> None:
|
||||
"""批量处理记录类型的事件"""
|
||||
db = create_session()
|
||||
|
||||
try:
|
||||
# 准备批量记录数据
|
||||
records: List[Dict[str, Any]] = []
|
||||
message_ids: List[str] = []
|
||||
message_fields: List[Dict[str, Any]] = []
|
||||
records: list[dict[str, Any]] = []
|
||||
message_ids: list[str] = []
|
||||
message_fields: list[dict[str, Any]] = []
|
||||
|
||||
for message_id, fields, event in messages:
|
||||
records.append(_event_to_record(event))
|
||||
@@ -305,7 +304,7 @@ class UsageQueueConsumer:
|
||||
db.rollback() # 清理批量失败的事务状态
|
||||
except Exception:
|
||||
pass
|
||||
success_ids: List[str] = []
|
||||
success_ids: list[str] = []
|
||||
for message_id, fields, event in messages:
|
||||
try:
|
||||
await self._apply_record_event(event, db=db)
|
||||
@@ -338,7 +337,7 @@ class UsageQueueConsumer:
|
||||
self,
|
||||
redis_client,
|
||||
message_id: str,
|
||||
fields: Dict[str, Any],
|
||||
fields: dict[str, Any],
|
||||
error: Exception,
|
||||
) -> None:
|
||||
retries = await self._get_delivery_count(redis_client, message_id)
|
||||
@@ -410,7 +409,7 @@ class UsageQueueConsumer:
|
||||
db.close()
|
||||
|
||||
async def _apply_record_event(
|
||||
self, event: UsageEvent, db: Optional[Session] = None
|
||||
self, event: UsageEvent, db: Session | None = None
|
||||
) -> None:
|
||||
"""处理记录类型事件(逐条写入,用于 fallback)
|
||||
|
||||
@@ -502,10 +501,10 @@ class UsageQueueConsumer:
|
||||
logger.debug(f"[usage-queue] metrics log failed: {exc}")
|
||||
|
||||
|
||||
_consumer_instance: Optional[UsageQueueConsumer] = None
|
||||
_consumer_instance: UsageQueueConsumer | None = None
|
||||
|
||||
|
||||
async def start_usage_queue_consumer() -> Optional[UsageQueueConsumer]:
|
||||
async def start_usage_queue_consumer() -> UsageQueueConsumer | None:
|
||||
global _consumer_instance
|
||||
if not config.usage_queue_enabled:
|
||||
return None
|
||||
|
||||
@@ -2,13 +2,13 @@
|
||||
Usage 事件定义与序列化工具(用于 Redis Streams)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from __future__ import annotations
|
||||
import json
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from typing import Any, Dict, Optional
|
||||
from typing import Any
|
||||
|
||||
USAGE_EVENT_VERSION = 1
|
||||
|
||||
@@ -34,7 +34,7 @@ def _sanitize_value(value: Any) -> Any:
|
||||
return str(value)
|
||||
|
||||
|
||||
def sanitize_payload(data: Dict[str, Any]) -> Dict[str, Any]:
|
||||
def sanitize_payload(data: dict[str, Any]) -> dict[str, Any]:
|
||||
return {str(k): _sanitize_value(v) for k, v in data.items()}
|
||||
|
||||
|
||||
@@ -43,9 +43,9 @@ class UsageEvent:
|
||||
event_type: UsageEventType
|
||||
request_id: str
|
||||
timestamp_ms: int
|
||||
data: Dict[str, Any]
|
||||
data: dict[str, Any]
|
||||
|
||||
def to_stream_fields(self) -> Dict[str, str]:
|
||||
def to_stream_fields(self) -> dict[str, str]:
|
||||
payload = {
|
||||
"v": USAGE_EVENT_VERSION,
|
||||
"type": self.event_type.value,
|
||||
@@ -56,7 +56,7 @@ class UsageEvent:
|
||||
return {"payload": json.dumps(payload, ensure_ascii=False)}
|
||||
|
||||
@classmethod
|
||||
def from_stream_fields(cls, fields: Dict[str, Any]) -> "UsageEvent":
|
||||
def from_stream_fields(cls, fields: dict[str, Any]) -> UsageEvent:
|
||||
raw = fields.get("payload")
|
||||
if not raw:
|
||||
raise ValueError("Missing payload field in usage event")
|
||||
@@ -76,8 +76,8 @@ def build_usage_event(
|
||||
*,
|
||||
event_type: UsageEventType,
|
||||
request_id: str,
|
||||
data: Dict[str, Any],
|
||||
timestamp_ms: Optional[int] = None,
|
||||
data: dict[str, Any],
|
||||
timestamp_ms: int | None = None,
|
||||
) -> UsageEvent:
|
||||
return UsageEvent(
|
||||
event_type=event_type,
|
||||
|
||||
@@ -22,7 +22,7 @@ await recorder.record(result)
|
||||
```
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, Optional
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
@@ -52,7 +52,7 @@ class UsageRecorder:
|
||||
user: User,
|
||||
api_key: ApiKey,
|
||||
client_ip: str = "unknown",
|
||||
request_id: Optional[str] = None,
|
||||
request_id: str | None = None,
|
||||
):
|
||||
self.db = db
|
||||
self.user = user
|
||||
@@ -75,8 +75,8 @@ class UsageRecorder:
|
||||
async def record_success(
|
||||
self,
|
||||
result: RequestResult,
|
||||
request_headers: Optional[Dict[str, str]] = None,
|
||||
request_body: Optional[Dict[str, Any]] = None,
|
||||
request_headers: dict[str, str] | None = None,
|
||||
request_body: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
记录成功请求的 Usage
|
||||
@@ -149,8 +149,8 @@ class UsageRecorder:
|
||||
async def record_failure(
|
||||
self,
|
||||
result: RequestResult,
|
||||
request_headers: Optional[Dict[str, str]] = None,
|
||||
request_body: Optional[Dict[str, Any]] = None,
|
||||
request_headers: dict[str, str] | None = None,
|
||||
request_body: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
记录失败请求的 Usage
|
||||
@@ -222,8 +222,8 @@ class UsageRecorder:
|
||||
model: str,
|
||||
response_time_ms: int,
|
||||
is_stream: bool = False,
|
||||
request_headers: Optional[Dict[str, str]] = None,
|
||||
request_body: Optional[Dict[str, Any]] = None,
|
||||
request_headers: dict[str, str] | None = None,
|
||||
request_body: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
从异常创建 RequestResult 并记录失败
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import func
|
||||
from sqlalchemy.orm import Session
|
||||
@@ -29,8 +29,8 @@ from src.services.system.config import SystemConfigService
|
||||
class UsageRecordParams:
|
||||
"""用量记录参数数据类,用于在内部方法间传递数据"""
|
||||
db: Session
|
||||
user: Optional[User]
|
||||
api_key: Optional[ApiKey]
|
||||
user: User | None
|
||||
api_key: ApiKey | None
|
||||
provider: str
|
||||
model: str
|
||||
input_tokens: int
|
||||
@@ -38,29 +38,29 @@ class UsageRecordParams:
|
||||
cache_creation_input_tokens: int
|
||||
cache_read_input_tokens: int
|
||||
request_type: str
|
||||
api_format: Optional[str]
|
||||
endpoint_api_format: Optional[str] # 端点原生 API 格式
|
||||
api_format: str | None
|
||||
endpoint_api_format: str | None # 端点原生 API 格式
|
||||
has_format_conversion: bool # 是否发生了格式转换
|
||||
is_stream: bool
|
||||
response_time_ms: Optional[int]
|
||||
first_byte_time_ms: Optional[int]
|
||||
response_time_ms: int | None
|
||||
first_byte_time_ms: int | None
|
||||
status_code: int
|
||||
error_message: Optional[str]
|
||||
metadata: Optional[Dict[str, Any]]
|
||||
request_headers: Optional[Dict[str, Any]]
|
||||
request_body: Optional[Any]
|
||||
provider_request_headers: Optional[Dict[str, Any]]
|
||||
response_headers: Optional[Dict[str, Any]]
|
||||
client_response_headers: Optional[Dict[str, Any]]
|
||||
response_body: Optional[Any]
|
||||
error_message: str | None
|
||||
metadata: dict[str, Any] | None
|
||||
request_headers: dict[str, Any] | None
|
||||
request_body: Any | None
|
||||
provider_request_headers: dict[str, Any] | None
|
||||
response_headers: dict[str, Any] | None
|
||||
client_response_headers: dict[str, Any] | None
|
||||
response_body: Any | None
|
||||
request_id: str
|
||||
provider_id: Optional[str]
|
||||
provider_endpoint_id: Optional[str]
|
||||
provider_api_key_id: Optional[str]
|
||||
provider_id: str | None
|
||||
provider_endpoint_id: str | None
|
||||
provider_api_key_id: str | None
|
||||
status: str
|
||||
cache_ttl_minutes: Optional[int]
|
||||
cache_ttl_minutes: int | None
|
||||
use_tiered_pricing: bool
|
||||
target_model: Optional[str]
|
||||
target_model: str | None
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
"""验证关键字段,确保数据完整性"""
|
||||
@@ -110,7 +110,7 @@ class UsageService:
|
||||
# ==================== 热力图缓存 ====================
|
||||
|
||||
@classmethod
|
||||
def _get_heatmap_cache_key(cls, user_id: Optional[str], include_actual_cost: bool) -> str:
|
||||
def _get_heatmap_cache_key(cls, user_id: str | None, include_actual_cost: bool) -> str:
|
||||
"""生成热力图缓存键"""
|
||||
cost_suffix = "with_cost" if include_actual_cost else "no_cost"
|
||||
if user_id:
|
||||
@@ -149,9 +149,9 @@ class UsageService:
|
||||
async def get_cached_heatmap(
|
||||
cls,
|
||||
db: Session,
|
||||
user_id: Optional[str] = None,
|
||||
user_id: str | None = None,
|
||||
include_actual_cost: bool = False,
|
||||
) -> Dict[str, Any]:
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
获取带缓存的热力图数据
|
||||
|
||||
@@ -220,8 +220,8 @@ class UsageService:
|
||||
def _build_usage_params(
|
||||
*,
|
||||
db: Session,
|
||||
user: Optional[User],
|
||||
api_key: Optional[ApiKey],
|
||||
user: User | None,
|
||||
api_key: ApiKey | None,
|
||||
provider: str,
|
||||
model: str,
|
||||
input_tokens: int,
|
||||
@@ -229,27 +229,27 @@ class UsageService:
|
||||
cache_creation_input_tokens: int,
|
||||
cache_read_input_tokens: int,
|
||||
request_type: str,
|
||||
api_format: Optional[str],
|
||||
endpoint_api_format: Optional[str],
|
||||
api_format: str | None,
|
||||
endpoint_api_format: str | None,
|
||||
has_format_conversion: bool,
|
||||
is_stream: bool,
|
||||
response_time_ms: Optional[int],
|
||||
first_byte_time_ms: Optional[int],
|
||||
response_time_ms: int | None,
|
||||
first_byte_time_ms: int | None,
|
||||
status_code: int,
|
||||
error_message: Optional[str],
|
||||
metadata: Optional[Dict[str, Any]],
|
||||
request_headers: Optional[Dict[str, Any]],
|
||||
request_body: Optional[Any],
|
||||
provider_request_headers: Optional[Dict[str, Any]],
|
||||
response_headers: Optional[Dict[str, Any]],
|
||||
client_response_headers: Optional[Dict[str, Any]],
|
||||
response_body: Optional[Any],
|
||||
error_message: str | None,
|
||||
metadata: dict[str, Any] | None,
|
||||
request_headers: dict[str, Any] | None,
|
||||
request_body: Any | None,
|
||||
provider_request_headers: dict[str, Any] | None,
|
||||
response_headers: dict[str, Any] | None,
|
||||
client_response_headers: dict[str, Any] | None,
|
||||
response_body: Any | None,
|
||||
request_id: str,
|
||||
provider_id: Optional[str],
|
||||
provider_endpoint_id: Optional[str],
|
||||
provider_api_key_id: Optional[str],
|
||||
provider_id: str | None,
|
||||
provider_endpoint_id: str | None,
|
||||
provider_api_key_id: str | None,
|
||||
status: str,
|
||||
target_model: Optional[str],
|
||||
target_model: str | None,
|
||||
# 成本计算结果
|
||||
input_cost: float,
|
||||
output_cost: float,
|
||||
@@ -261,13 +261,13 @@ class UsageService:
|
||||
# 价格信息
|
||||
input_price: float,
|
||||
output_price: float,
|
||||
cache_creation_price: Optional[float],
|
||||
cache_read_price: Optional[float],
|
||||
request_price: Optional[float],
|
||||
cache_creation_price: float | None,
|
||||
cache_read_price: float | None,
|
||||
request_price: float | None,
|
||||
# 倍率
|
||||
actual_rate_multiplier: float,
|
||||
is_free_tier: bool,
|
||||
) -> Dict[str, Any]:
|
||||
) -> dict[str, Any]:
|
||||
"""构建 Usage 记录的参数字典(内部方法,避免代码重复)"""
|
||||
|
||||
# 根据配置决定是否记录请求详情
|
||||
@@ -388,10 +388,10 @@ class UsageService:
|
||||
async def _get_rate_multiplier_and_free_tier(
|
||||
cls,
|
||||
db: Session,
|
||||
provider_api_key_id: Optional[str],
|
||||
provider_id: Optional[str],
|
||||
api_format: Optional[str] = None,
|
||||
) -> Tuple[float, bool]:
|
||||
provider_api_key_id: str | None,
|
||||
provider_id: str | None,
|
||||
api_format: str | None = None,
|
||||
) -> tuple[float, bool]:
|
||||
"""获取费率倍数和是否免费套餐(使用缓存)"""
|
||||
from src.services.cache.provider_cache import ProviderCacheService
|
||||
|
||||
@@ -409,12 +409,12 @@ class UsageService:
|
||||
output_tokens: int,
|
||||
cache_creation_input_tokens: int,
|
||||
cache_read_input_tokens: int,
|
||||
api_format: Optional[str],
|
||||
cache_ttl_minutes: Optional[int],
|
||||
api_format: str | None,
|
||||
cache_ttl_minutes: int | None,
|
||||
use_tiered_pricing: bool,
|
||||
is_failed_request: bool,
|
||||
) -> Tuple[float, float, float, float, float, float, float, float, float,
|
||||
Optional[float], Optional[float], Optional[float], Optional[int]]:
|
||||
) -> tuple[float, float, float, float, float, float, float, float, float,
|
||||
float | None, float | None, float | None, int | None]:
|
||||
"""计算所有成本相关数据
|
||||
|
||||
Returns:
|
||||
@@ -430,7 +430,7 @@ class UsageService:
|
||||
price_task = service.get_model_price_async(provider, model)
|
||||
request_price_task = service.get_request_price_async(provider, model)
|
||||
|
||||
tiered_pricing: Optional[dict] = None
|
||||
tiered_pricing: dict | None = None
|
||||
if use_tiered_pricing:
|
||||
tiered_pricing_task = service.get_tiered_pricing_async(provider, model)
|
||||
(input_price, output_price), request_price, tiered_pricing = await asyncio.gather(
|
||||
@@ -544,8 +544,8 @@ class UsageService:
|
||||
@staticmethod
|
||||
def _update_existing_usage(
|
||||
existing_usage: Usage,
|
||||
usage_params: Dict[str, Any],
|
||||
target_model: Optional[str],
|
||||
usage_params: dict[str, Any],
|
||||
target_model: str | None,
|
||||
) -> None:
|
||||
"""更新已存在的 Usage 记录(内部方法)"""
|
||||
# 更新关键字段
|
||||
@@ -630,7 +630,7 @@ class UsageService:
|
||||
@classmethod
|
||||
async def get_cache_prices_async(
|
||||
cls, db: Session, provider: str, model: str, input_price: float
|
||||
) -> Tuple[Optional[float], Optional[float]]:
|
||||
) -> tuple[float | None, float | None]:
|
||||
"""异步获取模型缓存价格(缓存创建价格,缓存读取价格)每1M tokens"""
|
||||
service = ModelCostService(db)
|
||||
return await service.get_cache_prices_async(provider, model, input_price)
|
||||
@@ -638,7 +638,7 @@ class UsageService:
|
||||
@classmethod
|
||||
def get_cache_prices(
|
||||
cls, db: Session, provider: str, model: str, input_price: float
|
||||
) -> Tuple[Optional[float], Optional[float]]:
|
||||
) -> tuple[float | None, float | None]:
|
||||
"""获取模型缓存价格(缓存创建价格,缓存读取价格)每1M tokens"""
|
||||
service = ModelCostService(db)
|
||||
return service.get_cache_prices(provider, model, input_price)
|
||||
@@ -646,13 +646,13 @@ class UsageService:
|
||||
@classmethod
|
||||
async def get_request_price_async(
|
||||
cls, db: Session, provider: str, model: str
|
||||
) -> Optional[float]:
|
||||
) -> float | None:
|
||||
"""异步获取模型按次计费价格"""
|
||||
service = ModelCostService(db)
|
||||
return await service.get_request_price_async(provider, model)
|
||||
|
||||
@classmethod
|
||||
def get_request_price(cls, db: Session, provider: str, model: str) -> Optional[float]:
|
||||
def get_request_price(cls, db: Session, provider: str, model: str) -> float | None:
|
||||
"""获取模型按次计费价格"""
|
||||
service = ModelCostService(db)
|
||||
return service.get_request_price(provider, model)
|
||||
@@ -665,9 +665,9 @@ class UsageService:
|
||||
output_price_per_1m: float,
|
||||
cache_creation_input_tokens: int = 0,
|
||||
cache_read_input_tokens: int = 0,
|
||||
cache_creation_price_per_1m: Optional[float] = None,
|
||||
cache_read_price_per_1m: Optional[float] = None,
|
||||
price_per_request: Optional[float] = None,
|
||||
cache_creation_price_per_1m: float | None = None,
|
||||
cache_read_price_per_1m: float | None = None,
|
||||
price_per_request: float | None = None,
|
||||
) -> tuple[float, float, float, float, float, float, float]:
|
||||
"""计算成本(价格是每百万tokens)- 固定价格模式
|
||||
|
||||
@@ -697,9 +697,9 @@ class UsageService:
|
||||
output_tokens: int,
|
||||
cache_creation_input_tokens: int = 0,
|
||||
cache_read_input_tokens: int = 0,
|
||||
api_format: Optional[str] = None,
|
||||
cache_ttl_minutes: Optional[int] = None,
|
||||
) -> tuple[float, float, float, float, float, float, float, Optional[int]]:
|
||||
api_format: str | None = None,
|
||||
cache_ttl_minutes: int | None = None,
|
||||
) -> tuple[float, float, float, float, float, float, float, int | None]:
|
||||
"""使用策略模式计算成本(支持阶梯计费)
|
||||
|
||||
根据 api_format 选择对应的计费策略,支持阶梯计费和 TTL 差异化。
|
||||
@@ -724,7 +724,7 @@ class UsageService:
|
||||
async def _prepare_usage_record(
|
||||
cls,
|
||||
params: UsageRecordParams,
|
||||
) -> Tuple[Dict[str, Any], float]:
|
||||
) -> tuple[dict[str, Any], float]:
|
||||
"""准备用量记录的共享逻辑
|
||||
|
||||
此方法提取了 record_usage 和 record_usage_async 的公共处理逻辑:
|
||||
@@ -817,8 +817,8 @@ class UsageService:
|
||||
@classmethod
|
||||
async def _prepare_usage_records_batch(
|
||||
cls,
|
||||
params_list: List[UsageRecordParams],
|
||||
) -> List[Tuple[Dict[str, Any], float, Optional[Exception]]]:
|
||||
params_list: list[UsageRecordParams],
|
||||
) -> list[tuple[dict[str, Any], float, Exception | None]]:
|
||||
"""批量并行准备用量记录(性能优化)
|
||||
|
||||
并行调用 _prepare_usage_record,提高批量处理效率。
|
||||
@@ -832,7 +832,7 @@ class UsageService:
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
async def prepare_single(params: UsageRecordParams) -> Tuple[Dict[str, Any], float, Optional[Exception]]:
|
||||
async def prepare_single(params: UsageRecordParams) -> tuple[dict[str, Any], float, Exception | None]:
|
||||
try:
|
||||
usage_params, total_cost = await cls._prepare_usage_record(params)
|
||||
return (usage_params, total_cost, None)
|
||||
@@ -845,7 +845,7 @@ class UsageService:
|
||||
# 避免一次性创建过多 task(并且 _prepare_usage_record 内部也可能包含并行调用)
|
||||
# 这里采用分批 gather 来限制并发量。
|
||||
chunk_size = 50
|
||||
results: List[Tuple[Dict[str, Any], float, Optional[Exception]]] = []
|
||||
results: list[tuple[dict[str, Any], float, Exception | None]] = []
|
||||
for i in range(0, len(params_list), chunk_size):
|
||||
chunk = params_list[i : i + chunk_size]
|
||||
chunk_results = await asyncio.gather(*(prepare_single(p) for p in chunk))
|
||||
@@ -856,8 +856,8 @@ class UsageService:
|
||||
async def record_usage_async(
|
||||
cls,
|
||||
db: Session,
|
||||
user: Optional[User],
|
||||
api_key: Optional[ApiKey],
|
||||
user: User | None,
|
||||
api_key: ApiKey | None,
|
||||
provider: str,
|
||||
model: str,
|
||||
input_tokens: int,
|
||||
@@ -865,29 +865,29 @@ class UsageService:
|
||||
cache_creation_input_tokens: int = 0,
|
||||
cache_read_input_tokens: int = 0,
|
||||
request_type: str = "chat",
|
||||
api_format: Optional[str] = None,
|
||||
endpoint_api_format: Optional[str] = None,
|
||||
api_format: str | None = None,
|
||||
endpoint_api_format: str | None = None,
|
||||
has_format_conversion: bool = False,
|
||||
is_stream: bool = False,
|
||||
response_time_ms: Optional[int] = None,
|
||||
first_byte_time_ms: Optional[int] = None,
|
||||
response_time_ms: int | None = None,
|
||||
first_byte_time_ms: int | None = None,
|
||||
status_code: int = 200,
|
||||
error_message: Optional[str] = None,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
request_headers: Optional[Dict[str, Any]] = None,
|
||||
request_body: Optional[Any] = None,
|
||||
provider_request_headers: Optional[Dict[str, Any]] = None,
|
||||
response_headers: Optional[Dict[str, Any]] = None,
|
||||
client_response_headers: Optional[Dict[str, Any]] = None,
|
||||
response_body: Optional[Any] = None,
|
||||
request_id: Optional[str] = None,
|
||||
provider_id: Optional[str] = None,
|
||||
provider_endpoint_id: Optional[str] = None,
|
||||
provider_api_key_id: Optional[str] = None,
|
||||
error_message: str | None = None,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
request_headers: dict[str, Any] | None = None,
|
||||
request_body: Any | None = None,
|
||||
provider_request_headers: dict[str, Any] | None = None,
|
||||
response_headers: dict[str, Any] | None = None,
|
||||
client_response_headers: dict[str, Any] | None = None,
|
||||
response_body: Any | None = None,
|
||||
request_id: str | None = None,
|
||||
provider_id: str | None = None,
|
||||
provider_endpoint_id: str | None = None,
|
||||
provider_api_key_id: str | None = None,
|
||||
status: str = "completed",
|
||||
cache_ttl_minutes: Optional[int] = None,
|
||||
cache_ttl_minutes: int | None = None,
|
||||
use_tiered_pricing: bool = True,
|
||||
target_model: Optional[str] = None,
|
||||
target_model: str | None = None,
|
||||
) -> Usage:
|
||||
"""异步记录使用量(简化版,仅插入新记录)
|
||||
|
||||
@@ -953,8 +953,8 @@ class UsageService:
|
||||
async def record_usage(
|
||||
cls,
|
||||
db: Session,
|
||||
user: Optional[User],
|
||||
api_key: Optional[ApiKey],
|
||||
user: User | None,
|
||||
api_key: ApiKey | None,
|
||||
provider: str,
|
||||
model: str,
|
||||
input_tokens: int,
|
||||
@@ -962,29 +962,29 @@ class UsageService:
|
||||
cache_creation_input_tokens: int = 0,
|
||||
cache_read_input_tokens: int = 0,
|
||||
request_type: str = "chat",
|
||||
api_format: Optional[str] = None,
|
||||
endpoint_api_format: Optional[str] = None,
|
||||
api_format: str | None = None,
|
||||
endpoint_api_format: str | None = None,
|
||||
has_format_conversion: bool = False,
|
||||
is_stream: bool = False,
|
||||
response_time_ms: Optional[int] = None,
|
||||
first_byte_time_ms: Optional[int] = None,
|
||||
response_time_ms: int | None = None,
|
||||
first_byte_time_ms: int | None = None,
|
||||
status_code: int = 200,
|
||||
error_message: Optional[str] = None,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
request_headers: Optional[Dict[str, Any]] = None,
|
||||
request_body: Optional[Any] = None,
|
||||
provider_request_headers: Optional[Dict[str, Any]] = None,
|
||||
response_headers: Optional[Dict[str, Any]] = None,
|
||||
client_response_headers: Optional[Dict[str, Any]] = None,
|
||||
response_body: Optional[Any] = None,
|
||||
request_id: Optional[str] = None,
|
||||
provider_id: Optional[str] = None,
|
||||
provider_endpoint_id: Optional[str] = None,
|
||||
provider_api_key_id: Optional[str] = None,
|
||||
error_message: str | None = None,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
request_headers: dict[str, Any] | None = None,
|
||||
request_body: Any | None = None,
|
||||
provider_request_headers: dict[str, Any] | None = None,
|
||||
response_headers: dict[str, Any] | None = None,
|
||||
client_response_headers: dict[str, Any] | None = None,
|
||||
response_body: Any | None = None,
|
||||
request_id: str | None = None,
|
||||
provider_id: str | None = None,
|
||||
provider_endpoint_id: str | None = None,
|
||||
provider_api_key_id: str | None = None,
|
||||
status: str = "completed",
|
||||
cache_ttl_minutes: Optional[int] = None,
|
||||
cache_ttl_minutes: int | None = None,
|
||||
use_tiered_pricing: bool = True,
|
||||
target_model: Optional[str] = None,
|
||||
target_model: str | None = None,
|
||||
) -> Usage:
|
||||
"""记录使用量(完整版,支持更新已存在记录和用户统计)
|
||||
|
||||
@@ -1113,8 +1113,8 @@ class UsageService:
|
||||
async def record_usage_batch(
|
||||
cls,
|
||||
db: Session,
|
||||
records: List[Dict[str, Any]],
|
||||
) -> List[Usage]:
|
||||
records: list[dict[str, Any]],
|
||||
) -> list[Usage]:
|
||||
"""批量记录使用量(高性能版,单次提交多条记录)
|
||||
|
||||
此方法针对高并发场景优化,特点:
|
||||
@@ -1139,9 +1139,9 @@ class UsageService:
|
||||
|
||||
# 分离需要更新和需要新建的记录
|
||||
request_ids = [r.get("request_id") for r in records if r.get("request_id")]
|
||||
existing_usages: Dict[str, Usage] = {}
|
||||
records_to_update: List[Dict[str, Any]] = []
|
||||
records_to_insert: List[Dict[str, Any]] = []
|
||||
existing_usages: dict[str, Usage] = {}
|
||||
records_to_update: list[dict[str, Any]] = []
|
||||
records_to_insert: list[dict[str, Any]] = []
|
||||
|
||||
if request_ids:
|
||||
# 查询已存在的 Usage 记录(包括 pending/streaming 状态)
|
||||
@@ -1174,13 +1174,13 @@ class UsageService:
|
||||
f"批量记录: 需要更新 {len(records_to_update)} 条已存在的 pending/streaming 记录"
|
||||
)
|
||||
|
||||
usages: List[Usage] = []
|
||||
user_costs: Dict[str, float] = defaultdict(float) # user_id -> total_cost
|
||||
apikey_stats: Dict[str, Dict[str, Any]] = defaultdict(
|
||||
usages: list[Usage] = []
|
||||
user_costs: dict[str, float] = defaultdict(float) # user_id -> total_cost
|
||||
apikey_stats: dict[str, dict[str, Any]] = defaultdict(
|
||||
lambda: {"requests": 0, "cost": 0.0, "is_standalone": False}
|
||||
)
|
||||
model_counts: Dict[str, int] = defaultdict(int) # model -> count
|
||||
provider_costs: Dict[str, float] = defaultdict(float) # provider_id -> cost
|
||||
model_counts: dict[str, int] = defaultdict(int) # model -> count
|
||||
provider_costs: dict[str, float] = defaultdict(float) # provider_id -> cost
|
||||
|
||||
# 合并所有需要处理的记录(用于预取 user/api_key)
|
||||
all_records = records_to_insert + records_to_update
|
||||
@@ -1189,12 +1189,12 @@ class UsageService:
|
||||
user_ids = {r.get("user_id") for r in all_records if r.get("user_id")}
|
||||
api_key_ids = {r.get("api_key_id") for r in all_records if r.get("api_key_id")}
|
||||
|
||||
users_map: Dict[str, User] = {}
|
||||
users_map: dict[str, User] = {}
|
||||
if user_ids:
|
||||
users = db.query(User).filter(User.id.in_(user_ids)).all()
|
||||
users_map = {str(u.id): u for u in users}
|
||||
|
||||
api_keys_map: Dict[str, ApiKey] = {}
|
||||
api_keys_map: dict[str, ApiKey] = {}
|
||||
if api_key_ids:
|
||||
api_keys = db.query(ApiKey).filter(ApiKey.id.in_(api_key_ids)).all()
|
||||
api_keys_map = {str(k.id): k for k in api_keys}
|
||||
@@ -1204,7 +1204,7 @@ class UsageService:
|
||||
total_count = len(all_records)
|
||||
|
||||
# 辅助函数:构建 UsageRecordParams
|
||||
def build_params(record: Dict[str, Any], request_id: str) -> UsageRecordParams:
|
||||
def build_params(record: dict[str, Any], request_id: str) -> UsageRecordParams:
|
||||
user_id = record.get("user_id")
|
||||
api_key_id = record.get("api_key_id")
|
||||
user = users_map.get(str(user_id)) if user_id else None
|
||||
@@ -1247,7 +1247,7 @@ class UsageService:
|
||||
)
|
||||
|
||||
# 构建所有参数并并行准备
|
||||
update_params_list: List[Tuple[Dict[str, Any], str, UsageRecordParams]] = []
|
||||
update_params_list: list[tuple[dict[str, Any], str, UsageRecordParams]] = []
|
||||
for record in records_to_update:
|
||||
request_id = record.get("request_id")
|
||||
if request_id and request_id in existing_usages:
|
||||
@@ -1260,7 +1260,7 @@ class UsageService:
|
||||
skipped_count += 1
|
||||
logger.warning("批量记录中参数构建失败: %s, request_id=%s", e, request_id)
|
||||
|
||||
insert_params_list: List[Tuple[Dict[str, Any], str, UsageRecordParams]] = []
|
||||
insert_params_list: list[tuple[dict[str, Any], str, UsageRecordParams]] = []
|
||||
for record in records_to_insert:
|
||||
request_id = record.get("request_id") or str(uuid.uuid4())[:8]
|
||||
try:
|
||||
@@ -1454,7 +1454,7 @@ class UsageService:
|
||||
user: User,
|
||||
estimated_tokens: int = 0,
|
||||
estimated_cost: float = 0,
|
||||
api_key: Optional[ApiKey] = None,
|
||||
api_key: ApiKey | None = None,
|
||||
) -> tuple[bool, str]:
|
||||
"""检查用户配额或独立Key余额
|
||||
|
||||
@@ -1513,12 +1513,12 @@ class UsageService:
|
||||
@staticmethod
|
||||
def get_usage_summary(
|
||||
db: Session,
|
||||
user_id: Optional[str] = None,
|
||||
api_key_id: Optional[str] = None,
|
||||
start_date: Optional[datetime] = None,
|
||||
end_date: Optional[datetime] = None,
|
||||
user_id: str | None = None,
|
||||
api_key_id: str | None = None,
|
||||
start_date: datetime | None = None,
|
||||
end_date: datetime | None = None,
|
||||
group_by: str = "day", # day, week, month
|
||||
) -> List[Dict[str, Any]]:
|
||||
) -> list[dict[str, Any]]:
|
||||
"""获取使用汇总"""
|
||||
|
||||
query = db.query(Usage)
|
||||
@@ -1596,12 +1596,12 @@ class UsageService:
|
||||
@staticmethod
|
||||
def get_daily_activity(
|
||||
db: Session,
|
||||
user_id: Optional[int] = None,
|
||||
start_date: Optional[datetime] = None,
|
||||
end_date: Optional[datetime] = None,
|
||||
user_id: int | None = None,
|
||||
start_date: datetime | None = None,
|
||||
end_date: datetime | None = None,
|
||||
window_days: int = 365,
|
||||
include_actual_cost: bool = False,
|
||||
) -> Dict[str, Any]:
|
||||
) -> dict[str, Any]:
|
||||
"""按天统计请求活跃度,用于渲染热力图。
|
||||
|
||||
优化策略:
|
||||
@@ -1627,7 +1627,7 @@ class UsageService:
|
||||
|
||||
today = now.date()
|
||||
today_start_dt = datetime.combine(today, datetime.min.time(), tzinfo=timezone.utc)
|
||||
aggregated: Dict[str, Dict[str, Any]] = {}
|
||||
aggregated: dict[str, dict[str, Any]] = {}
|
||||
|
||||
# 1. 从预计算表读取历史数据(不包括今天)
|
||||
if user_id:
|
||||
@@ -1724,7 +1724,7 @@ class UsageService:
|
||||
)
|
||||
|
||||
# 3. 构建返回结果
|
||||
days: List[Dict[str, Any]] = []
|
||||
days: list[dict[str, Any]] = []
|
||||
cursor = start_dt.date()
|
||||
end_date_only = end_dt.date()
|
||||
max_requests = 0
|
||||
@@ -1736,7 +1736,7 @@ class UsageService:
|
||||
total_tokens = stats.get("total_tokens", 0)
|
||||
total_cost = stats.get("total_cost_usd", 0.0)
|
||||
|
||||
entry: Dict[str, Any] = {
|
||||
entry: dict[str, Any] = {
|
||||
"date": iso_date,
|
||||
"requests": requests,
|
||||
"total_tokens": total_tokens,
|
||||
@@ -1762,10 +1762,10 @@ class UsageService:
|
||||
def get_top_users(
|
||||
db: Session,
|
||||
limit: int = 10,
|
||||
start_date: Optional[datetime] = None,
|
||||
end_date: Optional[datetime] = None,
|
||||
start_date: datetime | None = None,
|
||||
end_date: datetime | None = None,
|
||||
order_by: str = "cost", # cost, tokens, requests
|
||||
) -> List[Dict[str, Any]]:
|
||||
) -> list[dict[str, Any]]:
|
||||
"""获取使用量最高的用户"""
|
||||
|
||||
query = (
|
||||
@@ -1861,13 +1861,13 @@ class UsageService:
|
||||
cls,
|
||||
db: Session,
|
||||
request_id: str,
|
||||
user: Optional[User],
|
||||
api_key: Optional[ApiKey],
|
||||
user: User | None,
|
||||
api_key: ApiKey | None,
|
||||
model: str,
|
||||
is_stream: bool = False,
|
||||
api_format: Optional[str] = None,
|
||||
request_headers: Optional[Dict[str, Any]] = None,
|
||||
request_body: Optional[Any] = None,
|
||||
api_format: str | None = None,
|
||||
request_headers: dict[str, Any] | None = None,
|
||||
request_body: Any | None = None,
|
||||
) -> Usage:
|
||||
"""
|
||||
创建 pending 状态的使用记录(在请求开始时调用)
|
||||
@@ -1935,17 +1935,17 @@ class UsageService:
|
||||
db: Session,
|
||||
request_id: str,
|
||||
status: str,
|
||||
error_message: Optional[str] = None,
|
||||
provider: Optional[str] = None,
|
||||
target_model: Optional[str] = None,
|
||||
first_byte_time_ms: Optional[int] = None,
|
||||
provider_id: Optional[str] = None,
|
||||
provider_endpoint_id: Optional[str] = None,
|
||||
provider_api_key_id: Optional[str] = None,
|
||||
api_format: Optional[str] = None,
|
||||
endpoint_api_format: Optional[str] = None,
|
||||
has_format_conversion: Optional[bool] = None,
|
||||
) -> Optional[Usage]:
|
||||
error_message: str | None = None,
|
||||
provider: str | None = None,
|
||||
target_model: str | None = None,
|
||||
first_byte_time_ms: int | None = None,
|
||||
provider_id: str | None = None,
|
||||
provider_endpoint_id: str | None = None,
|
||||
provider_api_key_id: str | None = None,
|
||||
api_format: str | None = None,
|
||||
endpoint_api_format: str | None = None,
|
||||
has_format_conversion: bool | None = None,
|
||||
) -> Usage | None:
|
||||
"""
|
||||
快速更新使用记录状态
|
||||
|
||||
@@ -2016,8 +2016,8 @@ class UsageService:
|
||||
def _get_rate_multiplier_sync(
|
||||
db: Session,
|
||||
provider_api_key_id: str,
|
||||
api_format: Optional[str] = None,
|
||||
) -> Optional[float]:
|
||||
api_format: str | None = None,
|
||||
) -> float | None:
|
||||
"""
|
||||
同步获取 ProviderAPIKey 的 rate_multiplier
|
||||
|
||||
@@ -2048,9 +2048,9 @@ class UsageService:
|
||||
def get_active_requests(
|
||||
cls,
|
||||
db: Session,
|
||||
user_id: Optional[str] = None,
|
||||
user_id: str | None = None,
|
||||
limit: int = 50,
|
||||
) -> List[Usage]:
|
||||
) -> list[Usage]:
|
||||
"""
|
||||
获取活跃的请求(pending 或 streaming 状态)
|
||||
|
||||
@@ -2145,12 +2145,12 @@ class UsageService:
|
||||
def get_active_requests_status(
|
||||
cls,
|
||||
db: Session,
|
||||
ids: Optional[List[str]] = None,
|
||||
user_id: Optional[str] = None,
|
||||
ids: list[str] | None = None,
|
||||
user_id: str | None = None,
|
||||
default_timeout_seconds: int = 300,
|
||||
*,
|
||||
include_admin_fields: bool = False,
|
||||
) -> List[Dict[str, Any]]:
|
||||
) -> list[dict[str, Any]]:
|
||||
"""
|
||||
获取活跃请求状态(用于前端轮询),并自动清理超时的 pending/streaming 请求
|
||||
|
||||
@@ -2215,7 +2215,7 @@ class UsageService:
|
||||
|
||||
# 检查超时的 pending/streaming 请求
|
||||
# 收集可能超时的 usage_id 列表
|
||||
timeout_candidates: List[str] = []
|
||||
timeout_candidates: list[str] = []
|
||||
for r in records:
|
||||
if r.status in ("pending", "streaming") and r.created_at:
|
||||
# 使用全局配置的超时时间
|
||||
@@ -2312,7 +2312,7 @@ class UsageService:
|
||||
f"[Usage] 恢复 {len(completed_usage_ids)} 个已完成请求的状态(遥测回调丢失)"
|
||||
)
|
||||
|
||||
result: List[Dict[str, Any]] = []
|
||||
result: list[dict[str, Any]] = []
|
||||
for r in records:
|
||||
api_format = getattr(r, "api_format", None)
|
||||
endpoint_api_format = getattr(r, "endpoint_api_format", None)
|
||||
@@ -2326,7 +2326,7 @@ class UsageService:
|
||||
):
|
||||
has_format_conversion = not can_passthrough(api_format, endpoint_api_format)
|
||||
|
||||
item: Dict[str, Any] = {
|
||||
item: dict[str, Any] = {
|
||||
"id": r.id,
|
||||
"status": "failed" if r.id in timeout_ids else r.status,
|
||||
"input_tokens": r.input_tokens,
|
||||
@@ -2357,10 +2357,10 @@ class UsageService:
|
||||
@staticmethod
|
||||
def analyze_cache_affinity_ttl(
|
||||
db: Session,
|
||||
user_id: Optional[str] = None,
|
||||
api_key_id: Optional[str] = None,
|
||||
user_id: str | None = None,
|
||||
api_key_id: str | None = None,
|
||||
hours: int = 168,
|
||||
) -> Dict[str, Any]:
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
分析用户请求间隔分布,推荐合适的缓存亲和性 TTL
|
||||
|
||||
@@ -2437,7 +2437,7 @@ class UsageService:
|
||||
ORDER BY request_count DESC
|
||||
""")
|
||||
|
||||
params: Dict[str, Any] = {
|
||||
params: dict[str, Any] = {
|
||||
"start_date": start_date,
|
||||
}
|
||||
if user_id:
|
||||
@@ -2452,7 +2452,7 @@ class UsageService:
|
||||
group_ids = [row[0] for row in rows]
|
||||
|
||||
# 如果是按 user_id 分组,查询用户信息
|
||||
user_info_map: Dict[str, Dict[str, str]] = {}
|
||||
user_info_map: dict[str, dict[str, str]] = {}
|
||||
if group_by_field == "user_id" and group_ids:
|
||||
users = db.query(User).filter(User.id.in_(group_ids)).all()
|
||||
for user in users:
|
||||
@@ -2546,8 +2546,8 @@ class UsageService:
|
||||
|
||||
@staticmethod
|
||||
def _calculate_recommended_ttl(
|
||||
p75_interval: Optional[float],
|
||||
p90_interval: Optional[float],
|
||||
p75_interval: float | None,
|
||||
p90_interval: float | None,
|
||||
) -> int:
|
||||
"""
|
||||
根据请求间隔分布计算推荐的缓存 TTL
|
||||
@@ -2579,8 +2579,8 @@ class UsageService:
|
||||
@staticmethod
|
||||
def _get_ttl_recommendation_reason(
|
||||
ttl: int,
|
||||
p75_interval: Optional[float],
|
||||
p90_interval: Optional[float],
|
||||
p75_interval: float | None,
|
||||
p90_interval: float | None,
|
||||
) -> str:
|
||||
"""生成 TTL 推荐理由"""
|
||||
if p75_interval is None or p90_interval is None:
|
||||
@@ -2598,10 +2598,10 @@ class UsageService:
|
||||
@staticmethod
|
||||
def get_cache_hit_analysis(
|
||||
db: Session,
|
||||
user_id: Optional[str] = None,
|
||||
api_key_id: Optional[str] = None,
|
||||
user_id: str | None = None,
|
||||
api_key_id: str | None = None,
|
||||
hours: int = 168,
|
||||
) -> Dict[str, Any]:
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
分析缓存命中情况
|
||||
|
||||
@@ -2697,9 +2697,9 @@ class UsageService:
|
||||
db: Session,
|
||||
hours: int = 24,
|
||||
limit: int = 10000,
|
||||
user_id: Optional[str] = None,
|
||||
user_id: str | None = None,
|
||||
include_user_info: bool = False,
|
||||
) -> Dict[str, Any]:
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
获取请求间隔时间线数据,用于散点图展示
|
||||
|
||||
@@ -2808,7 +2808,7 @@ class UsageService:
|
||||
LIMIT :limit
|
||||
""")
|
||||
|
||||
params: Dict[str, Any] = {"start_date": start_date, "limit": limit}
|
||||
params: dict[str, Any] = {"start_date": start_date, "limit": limit}
|
||||
if user_id:
|
||||
params["user_id"] = user_id
|
||||
|
||||
@@ -2817,13 +2817,13 @@ class UsageService:
|
||||
|
||||
# 转换为时间线数据点
|
||||
points = []
|
||||
users_map: Dict[str, str] = {} # user_id -> username
|
||||
users_map: dict[str, str] = {} # user_id -> username
|
||||
models_set: set = set() # 收集所有出现的模型
|
||||
|
||||
if include_user_info and not user_id:
|
||||
for row in rows:
|
||||
created_at, row_user_id, model, username, interval_minutes = row
|
||||
point_data: Dict[str, Any] = {
|
||||
point_data: dict[str, Any] = {
|
||||
"x": created_at.isoformat(),
|
||||
"y": round(float(interval_minutes), 2),
|
||||
"user_id": str(row_user_id),
|
||||
@@ -2846,7 +2846,7 @@ class UsageService:
|
||||
models_set.add(model)
|
||||
points.append(point_data)
|
||||
|
||||
response: Dict[str, Any] = {
|
||||
response: dict[str, Any] = {
|
||||
"analysis_period_hours": hours,
|
||||
"total_points": len(points),
|
||||
"points": points,
|
||||
|
||||
@@ -5,8 +5,8 @@
|
||||
|
||||
import json
|
||||
import re
|
||||
import time
|
||||
from typing import Any, AsyncIterator, Dict, Optional, Tuple
|
||||
from typing import Any
|
||||
from collections.abc import AsyncIterator
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
@@ -30,19 +30,19 @@ class StreamUsageTracker:
|
||||
api_key: ApiKey,
|
||||
provider: str,
|
||||
model: str,
|
||||
request_headers: Optional[Dict[str, Any]] = None,
|
||||
provider_request_headers: Optional[Dict[str, Any]] = None,
|
||||
request_id: Optional[str] = None,
|
||||
start_time: Optional[float] = None,
|
||||
attempt_id: Optional[str] = None,
|
||||
request_headers: dict[str, Any] | None = None,
|
||||
provider_request_headers: dict[str, Any] | None = None,
|
||||
request_id: str | None = None,
|
||||
start_time: float | None = None,
|
||||
attempt_id: str | None = None,
|
||||
# Provider 侧追踪信息(用于记录真实成本)
|
||||
provider_id: Optional[str] = None,
|
||||
provider_endpoint_id: Optional[str] = None,
|
||||
provider_api_key_id: Optional[str] = None,
|
||||
provider_id: str | None = None,
|
||||
provider_endpoint_id: str | None = None,
|
||||
provider_api_key_id: str | None = None,
|
||||
# API 格式(用于选择正确的响应解析器)
|
||||
api_format: Optional[str] = None,
|
||||
api_format: str | None = None,
|
||||
# 格式转换信息
|
||||
endpoint_api_format: Optional[str] = None,
|
||||
endpoint_api_format: str | None = None,
|
||||
has_format_conversion: bool = False,
|
||||
):
|
||||
"""
|
||||
@@ -144,7 +144,7 @@ class StreamUsageTracker:
|
||||
self.error_message = error_message
|
||||
logger.debug(f"ID:{self.request_id} | 流式响应错误状态已设置 | 状态码:{status_code} | 错误:{error_message[:100]}")
|
||||
|
||||
def _update_complete_response(self, chunk: Dict[str, Any]):
|
||||
def _update_complete_response(self, chunk: dict[str, Any]):
|
||||
"""根据响应块更新完整响应结构"""
|
||||
try:
|
||||
# 更新响应ID
|
||||
@@ -245,7 +245,7 @@ class StreamUsageTracker:
|
||||
# 粗略估算:4个字符约等于1个token
|
||||
return max(1, total_chars // 4)
|
||||
|
||||
def _process_sse_event(self) -> Tuple[Optional[str], Optional[Dict[str, Any]]]:
|
||||
def _process_sse_event(self) -> tuple[str | None, dict[str, Any] | None]:
|
||||
"""
|
||||
处理缓冲区中的完整SSE事件
|
||||
|
||||
@@ -312,7 +312,7 @@ class StreamUsageTracker:
|
||||
|
||||
return content, usage
|
||||
|
||||
def parse_sse_line(self, line: str) -> Tuple[Optional[str], Optional[Dict[str, Any]]]:
|
||||
def parse_sse_line(self, line: str) -> tuple[str | None, dict[str, Any] | None]:
|
||||
"""
|
||||
解析单行SSE事件(使用统一响应解析器)
|
||||
|
||||
@@ -362,7 +362,7 @@ class StreamUsageTracker:
|
||||
|
||||
return content, usage
|
||||
|
||||
def parse_stream_chunk(self, chunk: bytes) -> Tuple[Optional[str], Optional[Dict[str, Any]]]:
|
||||
def parse_stream_chunk(self, chunk: bytes) -> tuple[str | None, dict[str, Any] | None]:
|
||||
"""
|
||||
解析流式响应块(处理原始字节流)
|
||||
|
||||
@@ -436,8 +436,8 @@ class StreamUsageTracker:
|
||||
async def track_stream(
|
||||
self,
|
||||
stream: AsyncIterator[str],
|
||||
request_data: Dict[str, Any],
|
||||
response_headers: Optional[Dict[str, Any]] = None,
|
||||
request_data: dict[str, Any],
|
||||
response_headers: dict[str, Any] | None = None,
|
||||
) -> AsyncIterator[str]:
|
||||
"""
|
||||
跟踪流式响应并计算用量
|
||||
@@ -663,7 +663,7 @@ class StreamUsageTracker:
|
||||
db_for_usage = self.db
|
||||
created_temp_session = False
|
||||
|
||||
def _load_user_and_key(db_session: Session) -> Tuple[Optional[User], Optional[ApiKey]]:
|
||||
def _load_user_and_key(db_session: Session) -> tuple[User | None, ApiKey | None]:
|
||||
local_user = (
|
||||
db_session.query(User).filter(User.id == self.user_id).first()
|
||||
if self.user_id
|
||||
@@ -801,19 +801,19 @@ class EnhancedStreamUsageTracker(StreamUsageTracker):
|
||||
api_key: ApiKey,
|
||||
provider: str,
|
||||
model: str,
|
||||
request_headers: Optional[Dict[str, Any]] = None,
|
||||
provider_request_headers: Optional[Dict[str, Any]] = None,
|
||||
request_id: Optional[str] = None,
|
||||
start_time: Optional[float] = None,
|
||||
attempt_id: Optional[str] = None,
|
||||
request_headers: dict[str, Any] | None = None,
|
||||
provider_request_headers: dict[str, Any] | None = None,
|
||||
request_id: str | None = None,
|
||||
start_time: float | None = None,
|
||||
attempt_id: str | None = None,
|
||||
# Provider 侧追踪信息(用于记录真实成本)
|
||||
provider_id: Optional[str] = None,
|
||||
provider_endpoint_id: Optional[str] = None,
|
||||
provider_api_key_id: Optional[str] = None,
|
||||
provider_id: str | None = None,
|
||||
provider_endpoint_id: str | None = None,
|
||||
provider_api_key_id: str | None = None,
|
||||
# API 格式(用于选择正确的响应解析器)
|
||||
api_format: Optional[str] = None,
|
||||
api_format: str | None = None,
|
||||
# 格式转换信息
|
||||
endpoint_api_format: Optional[str] = None,
|
||||
endpoint_api_format: str | None = None,
|
||||
has_format_conversion: bool = False,
|
||||
):
|
||||
super().__init__(
|
||||
@@ -907,8 +907,8 @@ class EnhancedStreamUsageTracker(StreamUsageTracker):
|
||||
async def track_stream(
|
||||
self,
|
||||
stream: AsyncIterator[str],
|
||||
request_data: Dict[str, Any],
|
||||
response_headers: Optional[Dict[str, Any]] = None,
|
||||
request_data: dict[str, Any],
|
||||
response_headers: dict[str, Any] | None = None,
|
||||
) -> AsyncIterator[str]:
|
||||
"""
|
||||
跟踪流式响应并更准确地计算用量
|
||||
@@ -1061,19 +1061,19 @@ def create_stream_tracker(
|
||||
provider: str,
|
||||
model: str,
|
||||
enhanced: bool = True,
|
||||
request_headers: Optional[Dict[str, Any]] = None,
|
||||
provider_request_headers: Optional[Dict[str, Any]] = None,
|
||||
request_id: Optional[str] = None,
|
||||
start_time: Optional[float] = None,
|
||||
attempt_id: Optional[str] = None,
|
||||
request_headers: dict[str, Any] | None = None,
|
||||
provider_request_headers: dict[str, Any] | None = None,
|
||||
request_id: str | None = None,
|
||||
start_time: float | None = None,
|
||||
attempt_id: str | None = None,
|
||||
# Provider 侧追踪信息(用于记录真实成本)
|
||||
provider_id: Optional[str] = None,
|
||||
provider_endpoint_id: Optional[str] = None,
|
||||
provider_api_key_id: Optional[str] = None,
|
||||
provider_id: str | None = None,
|
||||
provider_endpoint_id: str | None = None,
|
||||
provider_api_key_id: str | None = None,
|
||||
# API 格式(用于选择正确的响应解析器)
|
||||
api_format: Optional[str] = None,
|
||||
api_format: str | None = None,
|
||||
# 格式转换信息
|
||||
endpoint_api_format: Optional[str] = None,
|
||||
endpoint_api_format: str | None = None,
|
||||
has_format_conversion: bool = False,
|
||||
) -> StreamUsageTracker:
|
||||
"""
|
||||
|
||||
@@ -2,11 +2,10 @@
|
||||
Telemetry writer abstraction for stream usage.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any, Dict, Optional
|
||||
from typing import Any
|
||||
|
||||
from src.api.handlers.base.base_handler import MessageTelemetry
|
||||
from src.clients.redis_client import get_redis_client
|
||||
@@ -40,7 +39,7 @@ class DbTelemetryWriter(TelemetryWriter):
|
||||
def __init__(self, telemetry: MessageTelemetry) -> None:
|
||||
self._telemetry = telemetry
|
||||
|
||||
def _filter_kwargs(self, kwargs: Dict[str, Any]) -> Dict[str, Any]:
|
||||
def _filter_kwargs(self, kwargs: dict[str, Any]) -> dict[str, Any]:
|
||||
"""过滤掉 MessageTelemetry 不支持的参数"""
|
||||
return {k: v for k, v in kwargs.items() if k not in self._IGNORED_KWARGS}
|
||||
|
||||
@@ -101,7 +100,7 @@ class QueueTelemetryWriter(TelemetryWriter):
|
||||
logger.error(f"[usage-queue] XADD failed: {exc}")
|
||||
raise
|
||||
|
||||
def _truncate_body(self, value: Any) -> Optional[str]:
|
||||
def _truncate_body(self, value: Any) -> str | None:
|
||||
"""将 body 序列化为字符串,超长时截断并添加标记"""
|
||||
if value is None:
|
||||
return None
|
||||
@@ -116,9 +115,9 @@ class QueueTelemetryWriter(TelemetryWriter):
|
||||
raw = raw[:truncate_at] + "...[truncated]"
|
||||
return raw
|
||||
|
||||
def _build_event_data(self, **kwargs: Any) -> Dict[str, Any]:
|
||||
def _build_event_data(self, **kwargs: Any) -> dict[str, Any]:
|
||||
# 必需字段
|
||||
data: Dict[str, Any] = {
|
||||
data: dict[str, Any] = {
|
||||
"request_id": self.request_id,
|
||||
"user_id": self.user_id,
|
||||
"api_key_id": self.api_key_id,
|
||||
|
||||
@@ -3,14 +3,14 @@ API密钥管理服务
|
||||
"""
|
||||
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Any, Dict, List, Optional
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import func
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.crypto import crypto_service
|
||||
from src.core.logger import logger
|
||||
from src.models.database import ApiKey, Usage, User
|
||||
from src.models.database import ApiKey, Usage
|
||||
|
||||
|
||||
|
||||
@@ -21,15 +21,15 @@ class ApiKeyService:
|
||||
def create_api_key(
|
||||
db: Session,
|
||||
user_id: str, # UUID
|
||||
name: Optional[str] = None,
|
||||
allowed_providers: Optional[List[str]] = None,
|
||||
allowed_api_formats: Optional[List[str]] = None,
|
||||
allowed_models: Optional[List[str]] = None,
|
||||
rate_limit: Optional[int] = None,
|
||||
name: str | None = None,
|
||||
allowed_providers: list[str] | None = None,
|
||||
allowed_api_formats: list[str] | None = None,
|
||||
allowed_models: list[str] | None = None,
|
||||
rate_limit: int | None = None,
|
||||
concurrent_limit: int = 5,
|
||||
expire_days: Optional[int] = None,
|
||||
expires_at: Optional[datetime] = None, # 直接传入过期时间,优先于 expire_days
|
||||
initial_balance_usd: Optional[float] = None,
|
||||
expire_days: int | None = None,
|
||||
expires_at: datetime | None = None, # 直接传入过期时间,优先于 expire_days
|
||||
initial_balance_usd: float | None = None,
|
||||
is_standalone: bool = False,
|
||||
auto_delete_on_expiry: bool = False,
|
||||
) -> tuple[ApiKey, str]:
|
||||
@@ -89,20 +89,20 @@ class ApiKeyService:
|
||||
return api_key, key # 返回密钥对象和明文密钥
|
||||
|
||||
@staticmethod
|
||||
def get_api_key(db: Session, key_id: str) -> Optional[ApiKey]: # UUID
|
||||
def get_api_key(db: Session, key_id: str) -> ApiKey | None: # UUID
|
||||
"""获取API密钥"""
|
||||
return db.query(ApiKey).filter(ApiKey.id == key_id).first()
|
||||
|
||||
@staticmethod
|
||||
def get_api_key_by_key(db: Session, key: str) -> Optional[ApiKey]:
|
||||
def get_api_key_by_key(db: Session, key: str) -> ApiKey | None:
|
||||
"""通过密钥字符串获取API密钥"""
|
||||
key_hash = ApiKey.hash_key(key)
|
||||
return db.query(ApiKey).filter(ApiKey.key_hash == key_hash).first()
|
||||
|
||||
@staticmethod
|
||||
def list_user_api_keys(
|
||||
db: Session, user_id: str, is_active: Optional[bool] = None # UUID
|
||||
) -> List[ApiKey]:
|
||||
db: Session, user_id: str, is_active: bool | None = None # UUID
|
||||
) -> list[ApiKey]:
|
||||
"""列出用户的所有API密钥(不包括独立Key)"""
|
||||
query = db.query(ApiKey).filter(
|
||||
ApiKey.user_id == user_id, ApiKey.is_standalone == False # 排除独立Key
|
||||
@@ -114,7 +114,7 @@ class ApiKeyService:
|
||||
return query.order_by(ApiKey.created_at.desc()).all()
|
||||
|
||||
@staticmethod
|
||||
def list_standalone_api_keys(db: Session, is_active: Optional[bool] = None) -> List[ApiKey]:
|
||||
def list_standalone_api_keys(db: Session, is_active: bool | None = None) -> list[ApiKey]:
|
||||
"""列出所有独立余额Key(仅管理员可用)"""
|
||||
query = db.query(ApiKey).filter(ApiKey.is_standalone == True)
|
||||
|
||||
@@ -124,7 +124,7 @@ class ApiKeyService:
|
||||
return query.order_by(ApiKey.created_at.desc()).all()
|
||||
|
||||
@staticmethod
|
||||
def update_api_key(db: Session, key_id: str, **kwargs) -> Optional[ApiKey]: # UUID
|
||||
def update_api_key(db: Session, key_id: str, **kwargs) -> ApiKey | None: # UUID
|
||||
"""更新API密钥"""
|
||||
api_key = db.query(ApiKey).filter(ApiKey.id == key_id).first()
|
||||
if not api_key:
|
||||
@@ -186,7 +186,7 @@ class ApiKeyService:
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def get_remaining_balance(api_key: ApiKey) -> Optional[float]:
|
||||
def get_remaining_balance(api_key: ApiKey) -> float | None:
|
||||
"""计算剩余余额(仅用于独立Key)
|
||||
|
||||
Returns:
|
||||
@@ -203,7 +203,7 @@ class ApiKeyService:
|
||||
return max(0, remaining) # 不能为负数
|
||||
|
||||
@staticmethod
|
||||
def check_balance(api_key: ApiKey) -> tuple[bool, Optional[float]]:
|
||||
def check_balance(api_key: ApiKey) -> tuple[bool, float | None]:
|
||||
"""检查余额限制(仅用于独立Key)
|
||||
|
||||
Returns:
|
||||
@@ -259,7 +259,7 @@ class ApiKeyService:
|
||||
return is_allowed, api_key.rate_limit - request_count
|
||||
|
||||
@staticmethod
|
||||
def add_balance(db: Session, key_id: str, amount_usd: float) -> Optional[ApiKey]:
|
||||
def add_balance(db: Session, key_id: str, amount_usd: float) -> ApiKey | None:
|
||||
"""为独立余额Key调整余额
|
||||
|
||||
Args:
|
||||
@@ -355,9 +355,9 @@ class ApiKeyService:
|
||||
def get_api_key_stats(
|
||||
db: Session,
|
||||
key_id: str, # UUID
|
||||
start_date: Optional[datetime] = None,
|
||||
end_date: Optional[datetime] = None,
|
||||
) -> Dict[str, Any]:
|
||||
start_date: datetime | None = None,
|
||||
end_date: datetime | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""获取API密钥使用统计"""
|
||||
|
||||
api_key = db.query(ApiKey).filter(ApiKey.id == key_id).first()
|
||||
|
||||
@@ -2,7 +2,6 @@
|
||||
用户偏好设置服务
|
||||
"""
|
||||
|
||||
from typing import Optional
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
@@ -42,15 +41,15 @@ class PreferenceService:
|
||||
def update_preferences(
|
||||
db: Session,
|
||||
user_id: str, # UUID
|
||||
avatar_url: Optional[str] = None,
|
||||
bio: Optional[str] = None,
|
||||
default_provider_id: Optional[str] = None, # UUID
|
||||
theme: Optional[str] = None,
|
||||
language: Optional[str] = None,
|
||||
timezone: Optional[str] = None,
|
||||
email_notifications: Optional[bool] = None,
|
||||
usage_alerts: Optional[bool] = None,
|
||||
announcement_notifications: Optional[bool] = None,
|
||||
avatar_url: str | None = None,
|
||||
bio: str | None = None,
|
||||
default_provider_id: str | None = None, # UUID
|
||||
theme: str | None = None,
|
||||
language: str | None = None,
|
||||
timezone: str | None = None,
|
||||
email_notifications: bool | None = None,
|
||||
usage_alerts: bool | None = None,
|
||||
announcement_notifications: bool | None = None,
|
||||
) -> UserPreference:
|
||||
"""更新用户偏好设置"""
|
||||
preferences = PreferenceService.get_or_create_preferences(db, user_id)
|
||||
|
||||
@@ -3,8 +3,8 @@
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Any, Dict, List, Optional
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import and_, func
|
||||
from sqlalchemy.orm import Session
|
||||
@@ -25,15 +25,15 @@ class UserService:
|
||||
@retry_on_database_error(max_retries=3)
|
||||
def create_user(
|
||||
db: Session,
|
||||
email: Optional[str],
|
||||
email: str | None,
|
||||
username: str,
|
||||
password: str,
|
||||
role: UserRole = UserRole.USER,
|
||||
quota_usd: Optional[float] = 10.0,
|
||||
quota_usd: float | None = 10.0,
|
||||
email_verified: bool = False,
|
||||
allowed_providers: Optional[List[str]] = None,
|
||||
allowed_api_formats: Optional[List[str]] = None,
|
||||
allowed_models: Optional[List[str]] = None,
|
||||
allowed_providers: list[str] | None = None,
|
||||
allowed_api_formats: list[str] | None = None,
|
||||
allowed_models: list[str] | None = None,
|
||||
) -> User:
|
||||
"""创建新用户,quota_usd 为 None 表示无限制,email 为 None 表示无邮箱"""
|
||||
|
||||
@@ -90,7 +90,7 @@ class UserService:
|
||||
password: str,
|
||||
api_key_name: str = "默认密钥",
|
||||
role: UserRole = UserRole.USER,
|
||||
quota_usd: Optional[float] = 10.0,
|
||||
quota_usd: float | None = 10.0,
|
||||
concurrent_limit: int = 5,
|
||||
) -> tuple[User, ApiKey]:
|
||||
"""
|
||||
@@ -131,7 +131,7 @@ class UserService:
|
||||
return user, api_key, plain_key
|
||||
|
||||
@staticmethod
|
||||
def get_user(db: Session, user_id: str) -> Optional[User]:
|
||||
def get_user(db: Session, user_id: str) -> User | None:
|
||||
"""获取用户"""
|
||||
import random
|
||||
import time
|
||||
@@ -152,7 +152,7 @@ class UserService:
|
||||
raise e
|
||||
|
||||
@staticmethod
|
||||
def get_user_by_email(db: Session, email: str) -> Optional[User]:
|
||||
def get_user_by_email(db: Session, email: str) -> User | None:
|
||||
"""通过邮箱获取用户"""
|
||||
return db.query(User).filter(User.email == email).first()
|
||||
|
||||
@@ -161,9 +161,9 @@ class UserService:
|
||||
db: Session,
|
||||
skip: int = 0,
|
||||
limit: int = 100,
|
||||
role: Optional[UserRole] = None,
|
||||
is_active: Optional[bool] = None,
|
||||
) -> List[User]:
|
||||
role: UserRole | None = None,
|
||||
is_active: bool | None = None,
|
||||
) -> list[User]:
|
||||
"""列出用户"""
|
||||
query = db.query(User)
|
||||
|
||||
@@ -176,7 +176,7 @@ class UserService:
|
||||
|
||||
@staticmethod
|
||||
@transactional()
|
||||
def update_user(db: Session, user_id: str, **kwargs) -> Optional[User]:
|
||||
def update_user(db: Session, user_id: str, **kwargs) -> User | None:
|
||||
"""更新用户信息"""
|
||||
user = db.query(User).filter(User.id == user_id).first()
|
||||
if not user:
|
||||
@@ -322,8 +322,8 @@ class UserService:
|
||||
def update_user_quota(
|
||||
db: Session,
|
||||
user_id: str,
|
||||
quota_usd: Optional[float] = None,
|
||||
) -> Optional[User]:
|
||||
quota_usd: float | None = None,
|
||||
) -> User | None:
|
||||
"""更新用户配额"""
|
||||
user = db.query(User).filter(User.id == user_id).first()
|
||||
if not user:
|
||||
@@ -345,9 +345,9 @@ class UserService:
|
||||
def get_user_usage_stats(
|
||||
db: Session,
|
||||
user_id: str,
|
||||
start_date: Optional[datetime] = None,
|
||||
end_date: Optional[datetime] = None,
|
||||
) -> Dict[str, Any]:
|
||||
start_date: datetime | None = None,
|
||||
end_date: datetime | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""获取用户使用统计"""
|
||||
|
||||
query = db.query(Usage).filter(Usage.user_id == user_id)
|
||||
@@ -404,7 +404,7 @@ class UserService:
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def get_user_available_models(db: Session, user: User) -> List[Model]:
|
||||
def get_user_available_models(db: Session, user: User) -> list[Model]:
|
||||
"""获取用户可用的模型
|
||||
|
||||
通过 GlobalModel + Model 关联查询用户可用模型
|
||||
|
||||
Reference in New Issue
Block a user