mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 01:40:21 +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:
@@ -3,7 +3,6 @@ API Key认证插件
|
||||
支持从header中提取API Key进行认证
|
||||
"""
|
||||
|
||||
from typing import Optional
|
||||
|
||||
from fastapi import Request
|
||||
from sqlalchemy.orm import Session
|
||||
@@ -25,7 +24,7 @@ class ApiKeyAuthPlugin(AuthPlugin):
|
||||
def __init__(self):
|
||||
super().__init__(name="api_key", priority=10)
|
||||
|
||||
def get_credentials(self, request: Request) -> Optional[str]:
|
||||
def get_credentials(self, request: Request) -> str | None:
|
||||
"""
|
||||
从请求头中提取API Key
|
||||
|
||||
@@ -45,7 +44,7 @@ class ApiKeyAuthPlugin(AuthPlugin):
|
||||
|
||||
return None
|
||||
|
||||
async def authenticate(self, request: Request, db: Session) -> Optional[AuthContext]:
|
||||
async def authenticate(self, request: Request, db: Session) -> AuthContext | None:
|
||||
"""
|
||||
使用API Key进行认证
|
||||
"""
|
||||
|
||||
@@ -3,14 +3,14 @@
|
||||
定义认证插件的接口和认证上下文
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from abc import abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, List, Optional
|
||||
from typing import Any
|
||||
|
||||
from fastapi import Request
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from ..common import BasePlugin, HealthStatus, PluginMetadata
|
||||
from ..common import BasePlugin
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -22,11 +22,11 @@ class AuthContext:
|
||||
|
||||
user_id: int
|
||||
user_name: str
|
||||
api_key_id: Optional[int] = None
|
||||
api_key_name: Optional[str] = None
|
||||
permissions: Dict[str, bool] = None
|
||||
quota_info: Dict[str, Any] = None
|
||||
metadata: Dict[str, Any] = None
|
||||
api_key_id: int | None = None
|
||||
api_key_name: str | None = None
|
||||
permissions: dict[str, bool] = None
|
||||
quota_info: dict[str, Any] = None
|
||||
metadata: dict[str, Any] = None
|
||||
|
||||
def __post_init__(self):
|
||||
if self.permissions is None:
|
||||
@@ -49,9 +49,9 @@ class AuthPlugin(BasePlugin):
|
||||
author: str = "Unknown",
|
||||
description: str = "",
|
||||
api_version: str = "1.0",
|
||||
dependencies: List[str] = None,
|
||||
provides: List[str] = None,
|
||||
config: Dict[str, Any] = None,
|
||||
dependencies: list[str] = None,
|
||||
provides: list[str] = None,
|
||||
config: dict[str, Any] = None,
|
||||
):
|
||||
"""
|
||||
初始化认证插件
|
||||
@@ -80,7 +80,7 @@ class AuthPlugin(BasePlugin):
|
||||
)
|
||||
|
||||
@abstractmethod
|
||||
async def authenticate(self, request: Request, db: Session) -> Optional[AuthContext]:
|
||||
async def authenticate(self, request: Request, db: Session) -> AuthContext | None:
|
||||
"""
|
||||
执行认证
|
||||
|
||||
@@ -94,7 +94,7 @@ class AuthPlugin(BasePlugin):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_credentials(self, request: Request) -> Optional[str]:
|
||||
def get_credentials(self, request: Request) -> str | None:
|
||||
"""
|
||||
从请求中提取认证凭据
|
||||
|
||||
|
||||
@@ -4,7 +4,6 @@ JWT认证插件
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
from typing import Optional
|
||||
|
||||
from fastapi import Request
|
||||
from sqlalchemy.orm import Session
|
||||
@@ -26,7 +25,7 @@ class JwtAuthPlugin(AuthPlugin):
|
||||
def __init__(self):
|
||||
super().__init__(name="jwt", priority=20) # 高优先级,优先于API Key
|
||||
|
||||
def get_credentials(self, request: Request) -> Optional[str]:
|
||||
def get_credentials(self, request: Request) -> str | None:
|
||||
"""
|
||||
从Authorization header中提取JWT token
|
||||
|
||||
@@ -37,7 +36,7 @@ class JwtAuthPlugin(AuthPlugin):
|
||||
return auth_header.replace("Bearer ", "")
|
||||
return None
|
||||
|
||||
async def authenticate(self, request: Request, db: Session) -> Optional[AuthContext]:
|
||||
async def authenticate(self, request: Request, db: Session) -> AuthContext | None:
|
||||
"""
|
||||
使用JWT token进行认证
|
||||
"""
|
||||
|
||||
25
src/plugins/cache/base.py
vendored
25
src/plugins/cache/base.py
vendored
@@ -5,11 +5,10 @@
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
from abc import ABC, abstractmethod
|
||||
from datetime import timedelta
|
||||
from typing import Any, Dict, List, Optional
|
||||
from abc import abstractmethod
|
||||
from typing import Any
|
||||
|
||||
from ..common import BasePlugin, HealthStatus, PluginMetadata
|
||||
from ..common import BasePlugin
|
||||
|
||||
|
||||
class CachePlugin(BasePlugin):
|
||||
@@ -26,9 +25,9 @@ class CachePlugin(BasePlugin):
|
||||
author: str = "Unknown",
|
||||
description: str = "",
|
||||
api_version: str = "1.0",
|
||||
dependencies: List[str] = None,
|
||||
provides: List[str] = None,
|
||||
config: Dict[str, Any] = None,
|
||||
dependencies: list[str] = None,
|
||||
provides: list[str] = None,
|
||||
config: dict[str, Any] = None,
|
||||
):
|
||||
"""
|
||||
初始化缓存插件
|
||||
@@ -59,7 +58,7 @@ class CachePlugin(BasePlugin):
|
||||
self.max_size = self.config.get("max_size", 1000) # 最大缓存项数
|
||||
|
||||
@abstractmethod
|
||||
async def get(self, key: str) -> Optional[Any]:
|
||||
async def get(self, key: str) -> Any | None:
|
||||
"""
|
||||
获取缓存值
|
||||
|
||||
@@ -72,7 +71,7 @@ class CachePlugin(BasePlugin):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def set(self, key: str, value: Any, ttl: Optional[int] = None) -> bool:
|
||||
async def set(self, key: str, value: Any, ttl: int | None = None) -> bool:
|
||||
"""
|
||||
设置缓存值
|
||||
|
||||
@@ -123,7 +122,7 @@ class CachePlugin(BasePlugin):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def get_many(self, keys: List[str]) -> Dict[str, Any]:
|
||||
async def get_many(self, keys: list[str]) -> dict[str, Any]:
|
||||
"""
|
||||
批量获取缓存值
|
||||
|
||||
@@ -136,7 +135,7 @@ class CachePlugin(BasePlugin):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def set_many(self, items: Dict[str, Any], ttl: Optional[int] = None) -> bool:
|
||||
async def set_many(self, items: dict[str, Any], ttl: int | None = None) -> bool:
|
||||
"""
|
||||
批量设置缓存值
|
||||
|
||||
@@ -150,7 +149,7 @@ class CachePlugin(BasePlugin):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def get_stats(self) -> Dict[str, Any]:
|
||||
async def get_stats(self) -> dict[str, Any]:
|
||||
"""
|
||||
获取缓存统计信息
|
||||
|
||||
@@ -206,7 +205,7 @@ class CachePlugin(BasePlugin):
|
||||
"""
|
||||
return json.loads(value)
|
||||
|
||||
def configure(self, config: Dict[str, Any]):
|
||||
def configure(self, config: dict[str, Any]):
|
||||
"""
|
||||
配置插件
|
||||
|
||||
|
||||
24
src/plugins/cache/memory.py
vendored
24
src/plugins/cache/memory.py
vendored
@@ -7,7 +7,7 @@ import asyncio
|
||||
import threading
|
||||
import time
|
||||
from collections import OrderedDict
|
||||
from typing import Any, Dict, List, Optional
|
||||
from typing import Any
|
||||
|
||||
from .base import CachePlugin
|
||||
|
||||
@@ -18,10 +18,10 @@ class MemoryCachePlugin(CachePlugin):
|
||||
使用OrderedDict实现LRU缓存
|
||||
"""
|
||||
|
||||
def __init__(self, name: str = "memory", config: Dict[str, Any] = None):
|
||||
def __init__(self, name: str = "memory", config: dict[str, Any] = None):
|
||||
super().__init__(name, config)
|
||||
self._cache: OrderedDict = OrderedDict()
|
||||
self._expiry: Dict[str, float] = {}
|
||||
self._expiry: dict[str, float] = {}
|
||||
self._lock = threading.RLock()
|
||||
self._hits = 0
|
||||
self._misses = 0
|
||||
@@ -46,8 +46,12 @@ class MemoryCachePlugin(CachePlugin):
|
||||
await asyncio.sleep(self._cleanup_interval)
|
||||
await self._cleanup_expired()
|
||||
|
||||
loop = asyncio.get_event_loop()
|
||||
self._cleanup_task = loop.create_task(cleanup_loop())
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
self._cleanup_task = loop.create_task(cleanup_loop())
|
||||
except RuntimeError:
|
||||
# 没有运行的事件循环,稍后再启动
|
||||
pass
|
||||
|
||||
async def _cleanup_expired(self):
|
||||
"""清理过期的缓存项"""
|
||||
@@ -73,7 +77,7 @@ class MemoryCachePlugin(CachePlugin):
|
||||
self._expiry.pop(key, None)
|
||||
self._evictions += 1
|
||||
|
||||
async def get(self, key: str) -> Optional[Any]:
|
||||
async def get(self, key: str) -> Any | None:
|
||||
"""获取缓存值"""
|
||||
with self._lock:
|
||||
# 检查是否过期
|
||||
@@ -95,7 +99,7 @@ class MemoryCachePlugin(CachePlugin):
|
||||
self._misses += 1
|
||||
return None
|
||||
|
||||
async def set(self, key: str, value: Any, ttl: Optional[int] = None) -> bool:
|
||||
async def set(self, key: str, value: Any, ttl: int | None = None) -> bool:
|
||||
"""设置缓存值"""
|
||||
with self._lock:
|
||||
# 检查大小限制
|
||||
@@ -146,7 +150,7 @@ class MemoryCachePlugin(CachePlugin):
|
||||
self._expiry.clear()
|
||||
return True
|
||||
|
||||
async def get_many(self, keys: List[str]) -> Dict[str, Any]:
|
||||
async def get_many(self, keys: list[str]) -> dict[str, Any]:
|
||||
"""批量获取缓存值"""
|
||||
result = {}
|
||||
for key in keys:
|
||||
@@ -155,7 +159,7 @@ class MemoryCachePlugin(CachePlugin):
|
||||
result[key] = value
|
||||
return result
|
||||
|
||||
async def set_many(self, items: Dict[str, Any], ttl: Optional[int] = None) -> bool:
|
||||
async def set_many(self, items: dict[str, Any], ttl: int | None = None) -> bool:
|
||||
"""批量设置缓存值"""
|
||||
success = True
|
||||
for key, value in items.items():
|
||||
@@ -163,7 +167,7 @@ class MemoryCachePlugin(CachePlugin):
|
||||
success = False
|
||||
return success
|
||||
|
||||
async def get_stats(self) -> Dict[str, Any]:
|
||||
async def get_stats(self) -> dict[str, Any]:
|
||||
"""获取缓存统计信息"""
|
||||
total_requests = self._hits + self._misses
|
||||
hit_rate = self._hits / total_requests if total_requests > 0 else 0
|
||||
|
||||
@@ -3,10 +3,10 @@
|
||||
包含所有插件类型共享的类和接口
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from abc import ABC
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from typing import Any, Dict, List, Optional
|
||||
from typing import Any
|
||||
|
||||
from src.core.logger import logger
|
||||
|
||||
@@ -28,8 +28,8 @@ class PluginMetadata:
|
||||
author: str = "Unknown"
|
||||
description: str = ""
|
||||
api_version: str = "1.0"
|
||||
dependencies: List[str] = None
|
||||
provides: List[str] = None
|
||||
dependencies: list[str] = None
|
||||
provides: list[str] = None
|
||||
|
||||
def __post_init__(self):
|
||||
if self.dependencies is None:
|
||||
@@ -52,9 +52,9 @@ class BasePlugin(ABC):
|
||||
author: str = "Unknown",
|
||||
description: str = "",
|
||||
api_version: str = "1.0",
|
||||
dependencies: List[str] = None,
|
||||
provides: List[str] = None,
|
||||
config: Dict[str, Any] = None,
|
||||
dependencies: list[str] = None,
|
||||
provides: list[str] = None,
|
||||
config: dict[str, Any] = None,
|
||||
):
|
||||
"""
|
||||
初始化插件
|
||||
@@ -153,7 +153,7 @@ class BasePlugin(ABC):
|
||||
HealthStatus.HEALTHY if (self._initialized and self.enabled) else HealthStatus.UNHEALTHY
|
||||
)
|
||||
|
||||
def configure(self, config: Dict[str, Any]):
|
||||
def configure(self, config: dict[str, Any]):
|
||||
"""
|
||||
配置插件
|
||||
|
||||
@@ -177,7 +177,7 @@ class BasePlugin(ABC):
|
||||
"""检查插件是否已初始化"""
|
||||
return self._initialized
|
||||
|
||||
def validate_dependencies(self, available_plugins: Dict[str, List[str]]) -> List[str]:
|
||||
def validate_dependencies(self, available_plugins: dict[str, list[str]]) -> list[str]:
|
||||
"""
|
||||
验证插件依赖是否满足
|
||||
|
||||
|
||||
@@ -3,10 +3,9 @@
|
||||
定义负载均衡策略的接口
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from abc import abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from typing import Any, Dict, List, Optional
|
||||
from typing import Any
|
||||
|
||||
from ..common import BasePlugin
|
||||
|
||||
@@ -20,8 +19,8 @@ class ProviderCandidate:
|
||||
provider: Any # Provider 对象
|
||||
priority: int = 0 # 优先级(数字越大优先级越高)
|
||||
weight: float = 1.0 # 权重(影响被选中的概率)
|
||||
model: Optional[Any] = None # Model 对象(如果需要模型信息)
|
||||
metadata: Optional[Dict[str, Any]] = None # 额外元数据
|
||||
model: Any | None = None # Model 对象(如果需要模型信息)
|
||||
metadata: dict[str, Any] | None = None # 额外元数据
|
||||
|
||||
def __post_init__(self):
|
||||
if self.metadata is None:
|
||||
@@ -37,7 +36,7 @@ class SelectionResult:
|
||||
provider: Any # 选中的提供商
|
||||
priority: int # 该提供商的优先级
|
||||
weight: float # 该提供商的权重
|
||||
selection_metadata: Optional[Dict[str, Any]] = None # 选择过程的元数据
|
||||
selection_metadata: dict[str, Any] | None = None # 选择过程的元数据
|
||||
|
||||
def __post_init__(self):
|
||||
if self.selection_metadata is None:
|
||||
@@ -58,9 +57,9 @@ class LoadBalancerStrategy(BasePlugin):
|
||||
author: str = "Unknown",
|
||||
description: str = "",
|
||||
api_version: str = "1.0",
|
||||
dependencies: List[str] = None,
|
||||
provides: List[str] = None,
|
||||
config: Dict[str, Any] = None,
|
||||
dependencies: list[str] = None,
|
||||
provides: list[str] = None,
|
||||
config: dict[str, Any] = None,
|
||||
):
|
||||
"""
|
||||
初始化负载均衡策略
|
||||
@@ -90,8 +89,8 @@ class LoadBalancerStrategy(BasePlugin):
|
||||
|
||||
@abstractmethod
|
||||
async def select(
|
||||
self, candidates: List[ProviderCandidate], context: Optional[Dict[str, Any]] = None
|
||||
) -> Optional[SelectionResult]:
|
||||
self, candidates: list[ProviderCandidate], context: dict[str, Any] | None = None
|
||||
) -> SelectionResult | None:
|
||||
"""
|
||||
从候选提供商中选择一个
|
||||
|
||||
@@ -105,7 +104,7 @@ class LoadBalancerStrategy(BasePlugin):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def get_stats(self) -> Dict[str, Any]:
|
||||
async def get_stats(self) -> dict[str, Any]:
|
||||
"""
|
||||
获取负载均衡统计信息
|
||||
|
||||
@@ -118,8 +117,8 @@ class LoadBalancerStrategy(BasePlugin):
|
||||
self,
|
||||
provider: Any,
|
||||
success: bool,
|
||||
response_time: Optional[float] = None,
|
||||
error: Optional[Exception] = None,
|
||||
response_time: float | None = None,
|
||||
error: Exception | None = None,
|
||||
):
|
||||
"""
|
||||
记录请求结果(用于动态调整策略)
|
||||
|
||||
@@ -21,7 +21,7 @@ WARNING: 多进程环境注意事项
|
||||
|
||||
import time
|
||||
from collections import defaultdict
|
||||
from typing import Any, Dict, List, Optional
|
||||
from typing import Any
|
||||
|
||||
from src.core.logger import logger
|
||||
|
||||
@@ -49,7 +49,7 @@ class StickyPriorityStrategy(LoadBalancerStrategy):
|
||||
详见模块文档说明。
|
||||
"""
|
||||
|
||||
def __init__(self, config: Dict[str, Any] = None):
|
||||
def __init__(self, config: dict[str, Any] = None):
|
||||
config = config or {} # 确保 config 不为 None
|
||||
super().__init__(
|
||||
name="sticky_priority",
|
||||
@@ -68,7 +68,7 @@ class StickyPriorityStrategy(LoadBalancerStrategy):
|
||||
self.enable_auto_recovery = config.get("enable_auto_recovery", True) # 是否自动恢复
|
||||
|
||||
# 提供商健康状态追踪 {provider_id: health_info}
|
||||
self._provider_health: Dict[str, Dict[str, Any]] = defaultdict(
|
||||
self._provider_health: dict[str, dict[str, Any]] = defaultdict(
|
||||
lambda: {
|
||||
"consecutive_failures": 0,
|
||||
"last_failure_time": None,
|
||||
@@ -80,7 +80,7 @@ class StickyPriorityStrategy(LoadBalancerStrategy):
|
||||
|
||||
# 当前粘性提供商缓存 {cache_key: provider_id}
|
||||
# cache_key 可以是 api_key_id 或者其他标识
|
||||
self._sticky_providers: Dict[str, str] = {}
|
||||
self._sticky_providers: dict[str, str] = {}
|
||||
|
||||
# 统计信息
|
||||
self._stats = {
|
||||
@@ -92,8 +92,8 @@ class StickyPriorityStrategy(LoadBalancerStrategy):
|
||||
}
|
||||
|
||||
async def select(
|
||||
self, candidates: List[ProviderCandidate], context: Optional[Dict[str, Any]] = None
|
||||
) -> Optional[SelectionResult]:
|
||||
self, candidates: list[ProviderCandidate], context: dict[str, Any] | None = None
|
||||
) -> SelectionResult | None:
|
||||
"""
|
||||
从候选提供商中选择一个
|
||||
|
||||
@@ -180,7 +180,7 @@ class StickyPriorityStrategy(LoadBalancerStrategy):
|
||||
},
|
||||
)
|
||||
|
||||
def _get_cache_key(self, context: Optional[Dict[str, Any]]) -> str:
|
||||
def _get_cache_key(self, context: dict[str, Any] | None) -> str:
|
||||
"""
|
||||
生成缓存键,用于识别同一请求源
|
||||
|
||||
@@ -204,10 +204,10 @@ class StickyPriorityStrategy(LoadBalancerStrategy):
|
||||
return "default"
|
||||
|
||||
def _group_by_priority(
|
||||
self, candidates: List[ProviderCandidate]
|
||||
) -> Dict[int, List[ProviderCandidate]]:
|
||||
self, candidates: list[ProviderCandidate]
|
||||
) -> dict[int, list[ProviderCandidate]]:
|
||||
"""按优先级分组候选提供商"""
|
||||
groups: Dict[int, List[ProviderCandidate]] = {}
|
||||
groups: dict[int, list[ProviderCandidate]] = {}
|
||||
for candidate in candidates:
|
||||
priority = candidate.priority
|
||||
if priority not in groups:
|
||||
@@ -217,9 +217,9 @@ class StickyPriorityStrategy(LoadBalancerStrategy):
|
||||
|
||||
def _determine_sticky_provider(
|
||||
self,
|
||||
candidates: List[ProviderCandidate],
|
||||
candidates: list[ProviderCandidate],
|
||||
cache_key: str,
|
||||
context: Optional[Dict[str, Any]] = None,
|
||||
context: dict[str, Any] | None = None,
|
||||
) -> ProviderCandidate:
|
||||
"""
|
||||
确定粘性提供商
|
||||
@@ -289,8 +289,8 @@ class StickyPriorityStrategy(LoadBalancerStrategy):
|
||||
return False
|
||||
|
||||
def _select_backup_provider(
|
||||
self, candidates: List[ProviderCandidate]
|
||||
) -> Optional[ProviderCandidate]:
|
||||
self, candidates: list[ProviderCandidate]
|
||||
) -> ProviderCandidate | None:
|
||||
"""
|
||||
从候选列表中选择健康的备用提供商
|
||||
|
||||
@@ -340,8 +340,8 @@ class StickyPriorityStrategy(LoadBalancerStrategy):
|
||||
self,
|
||||
provider: Any,
|
||||
success: bool,
|
||||
response_time: Optional[float] = None,
|
||||
error: Optional[Exception] = None,
|
||||
response_time: float | None = None,
|
||||
error: Exception | None = None,
|
||||
):
|
||||
"""
|
||||
记录请求结果,更新健康状态
|
||||
@@ -377,7 +377,7 @@ class StickyPriorityStrategy(LoadBalancerStrategy):
|
||||
else:
|
||||
logger.debug(f"Recorded failed result for provider {provider.name}")
|
||||
|
||||
async def get_stats(self) -> Dict[str, Any]:
|
||||
async def get_stats(self) -> dict[str, Any]:
|
||||
"""获取统计信息"""
|
||||
# 计算健康状态
|
||||
healthy_count = sum(1 for info in self._provider_health.values() if info["is_healthy"])
|
||||
@@ -434,7 +434,7 @@ class StickyPriorityStrategy(LoadBalancerStrategy):
|
||||
}
|
||||
logger.info(f"Reset health status for provider {provider_id}")
|
||||
|
||||
async def clear_sticky_cache(self, cache_key: Optional[str] = None):
|
||||
async def clear_sticky_cache(self, cache_key: str | None = None):
|
||||
"""
|
||||
清除粘性提供商缓存
|
||||
|
||||
|
||||
@@ -8,14 +8,14 @@ import importlib
|
||||
import inspect
|
||||
import threading
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional, Type
|
||||
from typing import Any
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.plugins.auth.base import AuthPlugin
|
||||
from src.plugins.cache.base import CachePlugin
|
||||
|
||||
# 移除审计插件 - 审计功能现在是核心服务,不再作为插件
|
||||
from src.plugins.common import BasePlugin, HealthStatus, PluginMetadata
|
||||
from src.plugins.common import BasePlugin, HealthStatus
|
||||
from src.plugins.load_balancer.base import LoadBalancerStrategy
|
||||
from src.plugins.monitor.base import MonitorPlugin
|
||||
from src.plugins.notification.base import NotificationPlugin
|
||||
@@ -45,7 +45,7 @@ class PluginManager:
|
||||
# 移除 "audit" - 审计功能现在是核心服务
|
||||
}
|
||||
|
||||
def __init__(self, config: Optional[Dict[str, Any]] = None):
|
||||
def __init__(self, config: dict[str, Any] | None = None):
|
||||
"""
|
||||
初始化插件管理器
|
||||
|
||||
@@ -53,7 +53,7 @@ class PluginManager:
|
||||
config: 配置字典
|
||||
"""
|
||||
self.config = config or {}
|
||||
self.plugins: Dict[str, Dict[str, Any]] = {
|
||||
self.plugins: dict[str, dict[str, Any]] = {
|
||||
"auth": {},
|
||||
"rate_limit": {},
|
||||
"cache": {},
|
||||
@@ -63,7 +63,7 @@ class PluginManager:
|
||||
"load_balancer": {},
|
||||
# 移除 "audit" - 审计功能现在是核心服务
|
||||
}
|
||||
self.default_plugins: Dict[str, Optional[str]] = {
|
||||
self.default_plugins: dict[str, str | None] = {
|
||||
"auth": None,
|
||||
"rate_limit": None,
|
||||
"cache": None,
|
||||
@@ -74,7 +74,7 @@ class PluginManager:
|
||||
# 移除 "audit" - 审计功能现在是核心服务
|
||||
}
|
||||
# 跟踪因版本不兼容而跳过的插件
|
||||
self._incompatible_plugins: List[str] = []
|
||||
self._incompatible_plugins: list[str] = []
|
||||
|
||||
# 自动发现和加载插件
|
||||
self._auto_discover_plugins()
|
||||
@@ -210,7 +210,7 @@ class PluginManager:
|
||||
|
||||
logger.debug(f"Unregistered {plugin_type} plugin: {plugin_name}")
|
||||
|
||||
def get_plugin(self, plugin_type: str, plugin_name: Optional[str] = None) -> Optional[Any]:
|
||||
def get_plugin(self, plugin_type: str, plugin_name: str | None = None) -> Any | None:
|
||||
"""
|
||||
获取插件实例
|
||||
|
||||
@@ -238,7 +238,7 @@ class PluginManager:
|
||||
|
||||
return None
|
||||
|
||||
def get_plugins_by_type(self, plugin_type: str) -> List[Any]:
|
||||
def get_plugins_by_type(self, plugin_type: str) -> list[Any]:
|
||||
"""
|
||||
获取某个类型的所有插件
|
||||
|
||||
@@ -253,7 +253,7 @@ class PluginManager:
|
||||
|
||||
return list(self.plugins[plugin_type].values())
|
||||
|
||||
def get_enabled_plugins(self, plugin_type: str) -> List[Any]:
|
||||
def get_enabled_plugins(self, plugin_type: str) -> list[Any]:
|
||||
"""
|
||||
获取某个类型的所有启用的插件
|
||||
|
||||
@@ -302,7 +302,7 @@ class PluginManager:
|
||||
|
||||
return None
|
||||
|
||||
def get_stats(self) -> Dict[str, Any]:
|
||||
def get_stats(self) -> dict[str, Any]:
|
||||
"""
|
||||
获取插件管理器统计信息
|
||||
|
||||
@@ -340,7 +340,7 @@ class PluginManager:
|
||||
|
||||
return stats
|
||||
|
||||
async def initialize_all(self) -> Dict[str, bool]:
|
||||
async def initialize_all(self) -> dict[str, bool]:
|
||||
"""
|
||||
初始化所有插件
|
||||
|
||||
@@ -411,7 +411,7 @@ class PluginManager:
|
||||
except Exception as e:
|
||||
logger.error(f"Error during plugin shutdown: {e}")
|
||||
|
||||
def _sort_plugins_by_dependencies(self, plugins: List[BasePlugin]) -> List[BasePlugin]:
|
||||
def _sort_plugins_by_dependencies(self, plugins: list[BasePlugin]) -> list[BasePlugin]:
|
||||
"""
|
||||
按依赖关系对插件进行拓扑排序
|
||||
|
||||
@@ -466,7 +466,7 @@ class PluginManager:
|
||||
|
||||
return result
|
||||
|
||||
async def health_check_all(self) -> Dict[str, HealthStatus]:
|
||||
async def health_check_all(self) -> dict[str, HealthStatus]:
|
||||
"""
|
||||
检查所有插件的健康状态
|
||||
|
||||
@@ -496,7 +496,7 @@ class PluginManager:
|
||||
|
||||
return results
|
||||
|
||||
def validate_plugin_dependencies(self) -> Dict[str, List[str]]:
|
||||
def validate_plugin_dependencies(self) -> dict[str, list[str]]:
|
||||
"""
|
||||
验证所有插件的依赖关系
|
||||
|
||||
@@ -520,7 +520,7 @@ class PluginManager:
|
||||
return results
|
||||
|
||||
def reload_plugin_config(
|
||||
self, plugin_type: str, plugin_name: str, new_config: Dict[str, Any]
|
||||
self, plugin_type: str, plugin_name: str, new_config: dict[str, Any]
|
||||
) -> bool:
|
||||
"""
|
||||
重新加载插件配置
|
||||
@@ -547,11 +547,11 @@ class PluginManager:
|
||||
|
||||
|
||||
# 全局插件管理器实例
|
||||
_plugin_manager: Optional[PluginManager] = None
|
||||
_plugin_manager: PluginManager | None = None
|
||||
_plugin_manager_lock = threading.Lock()
|
||||
|
||||
|
||||
def get_plugin_manager(config: Optional[Dict[str, Any]] = None) -> PluginManager:
|
||||
def get_plugin_manager(config: dict[str, Any] | None = None) -> PluginManager:
|
||||
"""
|
||||
获取全局插件管理器实例(线程安全)
|
||||
|
||||
|
||||
@@ -6,7 +6,7 @@
|
||||
from abc import abstractmethod
|
||||
from datetime import datetime, timezone
|
||||
from enum import Enum
|
||||
from typing import Any, Dict, List, Optional
|
||||
from typing import Any
|
||||
|
||||
from src.plugins.common import BasePlugin
|
||||
|
||||
@@ -28,9 +28,9 @@ class Metric:
|
||||
name: str,
|
||||
value: float,
|
||||
metric_type: MetricType,
|
||||
labels: Optional[Dict[str, str]] = None,
|
||||
timestamp: Optional[datetime] = None,
|
||||
description: Optional[str] = None,
|
||||
labels: dict[str, str] | None = None,
|
||||
timestamp: datetime | None = None,
|
||||
description: str | None = None,
|
||||
):
|
||||
self.name = name
|
||||
self.value = value
|
||||
@@ -46,7 +46,7 @@ class MonitorPlugin(BasePlugin):
|
||||
所有监控插件必须继承此类并实现相关方法
|
||||
"""
|
||||
|
||||
def __init__(self, name: str, config: Dict[str, Any] = None):
|
||||
def __init__(self, name: str, config: dict[str, Any] = None):
|
||||
"""
|
||||
初始化监控插件
|
||||
|
||||
@@ -71,7 +71,7 @@ class MonitorPlugin(BasePlugin):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def record_batch(self, metrics: List[Metric]):
|
||||
async def record_batch(self, metrics: list[Metric]):
|
||||
"""
|
||||
批量记录指标
|
||||
|
||||
@@ -81,7 +81,7 @@ class MonitorPlugin(BasePlugin):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def increment(self, name: str, value: float = 1, labels: Optional[Dict[str, str]] = None):
|
||||
async def increment(self, name: str, value: float = 1, labels: dict[str, str] | None = None):
|
||||
"""
|
||||
增加计数器
|
||||
|
||||
@@ -93,7 +93,7 @@ class MonitorPlugin(BasePlugin):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def gauge(self, name: str, value: float, labels: Optional[Dict[str, str]] = None):
|
||||
async def gauge(self, name: str, value: float, labels: dict[str, str] | None = None):
|
||||
"""
|
||||
设置仪表值
|
||||
|
||||
@@ -109,8 +109,8 @@ class MonitorPlugin(BasePlugin):
|
||||
self,
|
||||
name: str,
|
||||
value: float,
|
||||
labels: Optional[Dict[str, str]] = None,
|
||||
buckets: Optional[List[float]] = None,
|
||||
labels: dict[str, str] | None = None,
|
||||
buckets: list[float] | None = None,
|
||||
):
|
||||
"""
|
||||
记录直方图数据
|
||||
@@ -124,7 +124,7 @@ class MonitorPlugin(BasePlugin):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def timing(self, name: str, duration: float, labels: Optional[Dict[str, str]] = None):
|
||||
async def timing(self, name: str, duration: float, labels: dict[str, str] | None = None):
|
||||
"""
|
||||
记录时间指标
|
||||
|
||||
@@ -143,7 +143,7 @@ class MonitorPlugin(BasePlugin):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def get_stats(self) -> Dict[str, Any]:
|
||||
async def get_stats(self) -> dict[str, Any]:
|
||||
"""
|
||||
获取插件统计信息
|
||||
|
||||
@@ -158,8 +158,8 @@ class MonitorPlugin(BasePlugin):
|
||||
endpoint: str,
|
||||
status_code: int,
|
||||
duration: float,
|
||||
provider: Optional[str] = None,
|
||||
model: Optional[str] = None,
|
||||
provider: str | None = None,
|
||||
model: str | None = None,
|
||||
):
|
||||
"""
|
||||
记录API请求指标(便捷方法)
|
||||
@@ -187,7 +187,10 @@ class MonitorPlugin(BasePlugin):
|
||||
# 异步记录指标
|
||||
import asyncio
|
||||
|
||||
loop = asyncio.get_event_loop()
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
return # 没有事件循环时跳过
|
||||
|
||||
# 请求计数
|
||||
loop.create_task(self.increment("http_requests_total", labels=labels))
|
||||
@@ -205,7 +208,7 @@ class MonitorPlugin(BasePlugin):
|
||||
model: str,
|
||||
input_tokens: int,
|
||||
output_tokens: int,
|
||||
cost: Optional[float] = None,
|
||||
cost: float | None = None,
|
||||
):
|
||||
"""
|
||||
记录Token使用指标(便捷方法)
|
||||
@@ -221,7 +224,10 @@ class MonitorPlugin(BasePlugin):
|
||||
|
||||
import asyncio
|
||||
|
||||
loop = asyncio.get_event_loop()
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
return # 没有事件循环时跳过
|
||||
|
||||
# Token计数
|
||||
loop.create_task(self.increment("tokens_input_total", input_tokens, labels=labels))
|
||||
@@ -234,7 +240,7 @@ class MonitorPlugin(BasePlugin):
|
||||
if cost is not None:
|
||||
loop.create_task(self.increment("usage_cost_total", cost, labels=labels))
|
||||
|
||||
def configure(self, config: Dict[str, Any]):
|
||||
def configure(self, config: dict[str, Any]):
|
||||
"""
|
||||
配置插件
|
||||
|
||||
|
||||
@@ -4,9 +4,7 @@ Prometheus监控插件
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from collections import defaultdict
|
||||
from typing import Any, Dict, List, Optional
|
||||
from typing import Any
|
||||
|
||||
try:
|
||||
from prometheus_client import REGISTRY, Counter, Gauge, Histogram, Summary, generate_latest
|
||||
@@ -28,7 +26,7 @@ class PrometheusPlugin(MonitorPlugin):
|
||||
使用prometheus_client库导出指标
|
||||
"""
|
||||
|
||||
def __init__(self, name: str = "prometheus", config: Dict[str, Any] = None):
|
||||
def __init__(self, name: str = "prometheus", config: dict[str, Any] = None):
|
||||
super().__init__(name, config)
|
||||
|
||||
# Check if prometheus_client is available
|
||||
@@ -38,10 +36,10 @@ class PrometheusPlugin(MonitorPlugin):
|
||||
return
|
||||
|
||||
# 指标注册表
|
||||
self._metrics: Dict[str, Any] = {}
|
||||
self._buffer: List[Metric] = []
|
||||
self._metrics: dict[str, Any] = {}
|
||||
self._buffer: list[Metric] = []
|
||||
self._lock = asyncio.Lock()
|
||||
self._flush_task: Optional[asyncio.Task] = None # 跟踪后台任务
|
||||
self._flush_task: asyncio.Task | None = None # 跟踪后台任务
|
||||
|
||||
# 预定义常用指标
|
||||
self._init_default_metrics()
|
||||
@@ -132,13 +130,13 @@ class PrometheusPlugin(MonitorPlugin):
|
||||
|
||||
# 保存任务句柄以便后续取消
|
||||
try:
|
||||
loop = asyncio.get_event_loop()
|
||||
loop = asyncio.get_running_loop()
|
||||
self._flush_task = loop.create_task(flush_loop())
|
||||
except RuntimeError:
|
||||
# 如果没有运行的事件循环,任务将在后续创建
|
||||
logger.warning("No event loop available for Prometheus flush task")
|
||||
|
||||
def _get_or_create_metric(self, name: str, metric_type: MetricType, labels: List[str] = None):
|
||||
def _get_or_create_metric(self, name: str, metric_type: MetricType, labels: list[str] = None):
|
||||
"""获取或创建指标"""
|
||||
if name not in self._metrics:
|
||||
labels = labels or []
|
||||
@@ -162,7 +160,7 @@ class PrometheusPlugin(MonitorPlugin):
|
||||
if len(self._buffer) >= self.batch_size:
|
||||
await self.flush()
|
||||
|
||||
async def record_batch(self, metrics: List[Metric]):
|
||||
async def record_batch(self, metrics: list[Metric]):
|
||||
"""批量记录指标"""
|
||||
async with self._lock:
|
||||
self._buffer.extend(metrics)
|
||||
@@ -171,7 +169,7 @@ class PrometheusPlugin(MonitorPlugin):
|
||||
if len(self._buffer) >= self.batch_size:
|
||||
await self.flush()
|
||||
|
||||
async def increment(self, name: str, value: float = 1, labels: Optional[Dict[str, str]] = None):
|
||||
async def increment(self, name: str, value: float = 1, labels: dict[str, str] | None = None):
|
||||
"""增加计数器"""
|
||||
try:
|
||||
if name in self._metrics:
|
||||
@@ -194,7 +192,7 @@ class PrometheusPlugin(MonitorPlugin):
|
||||
# 记录错误但不中断
|
||||
logger.warning(f"Error recording metric {name}: {e}")
|
||||
|
||||
async def gauge(self, name: str, value: float, labels: Optional[Dict[str, str]] = None):
|
||||
async def gauge(self, name: str, value: float, labels: dict[str, str] | None = None):
|
||||
"""设置仪表值"""
|
||||
try:
|
||||
if name in self._metrics:
|
||||
@@ -219,8 +217,8 @@ class PrometheusPlugin(MonitorPlugin):
|
||||
self,
|
||||
name: str,
|
||||
value: float,
|
||||
labels: Optional[Dict[str, str]] = None,
|
||||
buckets: Optional[List[float]] = None,
|
||||
labels: dict[str, str] | None = None,
|
||||
buckets: list[float] | None = None,
|
||||
):
|
||||
"""记录直方图数据"""
|
||||
try:
|
||||
@@ -249,7 +247,7 @@ class PrometheusPlugin(MonitorPlugin):
|
||||
except Exception as e:
|
||||
logger.warning(f"Error recording histogram {name}: {e}")
|
||||
|
||||
async def timing(self, name: str, duration: float, labels: Optional[Dict[str, str]] = None):
|
||||
async def timing(self, name: str, duration: float, labels: dict[str, str] | None = None):
|
||||
"""记录时间指标"""
|
||||
# 使用直方图记录时间
|
||||
await self.histogram(f"{name}_seconds", duration, labels)
|
||||
@@ -272,7 +270,7 @@ class PrometheusPlugin(MonitorPlugin):
|
||||
# 清空缓冲区
|
||||
self._buffer.clear()
|
||||
|
||||
async def get_stats(self) -> Dict[str, Any]:
|
||||
async def get_stats(self) -> dict[str, Any]:
|
||||
"""获取插件统计信息"""
|
||||
return {
|
||||
"type": "prometheus",
|
||||
|
||||
@@ -8,7 +8,7 @@ import json
|
||||
from abc import abstractmethod
|
||||
from datetime import datetime, timezone
|
||||
from enum import Enum
|
||||
from typing import Any, Dict, List, Optional
|
||||
from typing import Any
|
||||
|
||||
from src.plugins.common import BasePlugin
|
||||
|
||||
@@ -30,12 +30,12 @@ class Notification:
|
||||
title: str,
|
||||
message: str,
|
||||
level: NotificationLevel = NotificationLevel.INFO,
|
||||
notification_type: Optional[str] = None,
|
||||
source: Optional[str] = None,
|
||||
timestamp: Optional[datetime] = None,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
recipient: Optional[str] = None,
|
||||
tags: Optional[List[str]] = None,
|
||||
notification_type: str | None = None,
|
||||
source: str | None = None,
|
||||
timestamp: datetime | None = None,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
recipient: str | None = None,
|
||||
tags: list[str] | None = None,
|
||||
):
|
||||
self.title = title
|
||||
self.message = message
|
||||
@@ -47,7 +47,7 @@ class Notification:
|
||||
self.recipient = recipient
|
||||
self.tags = tags or []
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""转换为字典"""
|
||||
return {
|
||||
"title": self.title,
|
||||
@@ -65,7 +65,7 @@ class Notification:
|
||||
"""转换为JSON"""
|
||||
return json.dumps(self.to_dict(), default=str)
|
||||
|
||||
def format_message(self, template: Optional[str] = None) -> str:
|
||||
def format_message(self, template: str | None = None) -> str:
|
||||
"""格式化消息"""
|
||||
if template:
|
||||
return template.format(
|
||||
@@ -90,7 +90,7 @@ class NotificationPlugin(BasePlugin):
|
||||
提供统一的重试机制,子类只需实现 _do_send 和 _do_send_batch 方法
|
||||
"""
|
||||
|
||||
def __init__(self, name: str = "notification", config: Dict[str, Any] = None):
|
||||
def __init__(self, name: str = "notification", config: dict[str, Any] = None):
|
||||
# 调用父类初始化,设置metadata
|
||||
super().__init__(
|
||||
name=name, config=config, description="Notification Plugin", version="1.0.0"
|
||||
@@ -160,7 +160,7 @@ class NotificationPlugin(BasePlugin):
|
||||
"""
|
||||
pass
|
||||
|
||||
async def send_batch(self, notifications: List[Notification]) -> Dict[str, Any]:
|
||||
async def send_batch(self, notifications: list[Notification]) -> dict[str, Any]:
|
||||
"""
|
||||
批量发送通知(带重试机制)
|
||||
|
||||
@@ -210,7 +210,7 @@ class NotificationPlugin(BasePlugin):
|
||||
}
|
||||
|
||||
@abstractmethod
|
||||
async def _do_send_batch(self, notifications: List[Notification]) -> Dict[str, Any]:
|
||||
async def _do_send_batch(self, notifications: list[Notification]) -> dict[str, Any]:
|
||||
"""
|
||||
实际批量发送通知(子类实现)
|
||||
|
||||
@@ -237,8 +237,8 @@ class NotificationPlugin(BasePlugin):
|
||||
async def send_error(
|
||||
self,
|
||||
error: Exception,
|
||||
context: Optional[Dict[str, Any]] = None,
|
||||
recipient: Optional[str] = None,
|
||||
context: dict[str, Any] | None = None,
|
||||
recipient: str | None = None,
|
||||
) -> bool:
|
||||
"""发送错误通知"""
|
||||
notification = Notification(
|
||||
@@ -256,8 +256,8 @@ class NotificationPlugin(BasePlugin):
|
||||
self,
|
||||
title: str,
|
||||
message: str,
|
||||
context: Optional[Dict[str, Any]] = None,
|
||||
recipient: Optional[str] = None,
|
||||
context: dict[str, Any] | None = None,
|
||||
recipient: str | None = None,
|
||||
) -> bool:
|
||||
"""发送警告通知"""
|
||||
notification = Notification(
|
||||
@@ -275,8 +275,8 @@ class NotificationPlugin(BasePlugin):
|
||||
self,
|
||||
title: str,
|
||||
message: str,
|
||||
context: Optional[Dict[str, Any]] = None,
|
||||
recipient: Optional[str] = None,
|
||||
context: dict[str, Any] | None = None,
|
||||
recipient: str | None = None,
|
||||
) -> bool:
|
||||
"""发送信息通知"""
|
||||
notification = Notification(
|
||||
@@ -294,8 +294,8 @@ class NotificationPlugin(BasePlugin):
|
||||
self,
|
||||
title: str,
|
||||
message: str,
|
||||
context: Optional[Dict[str, Any]] = None,
|
||||
recipient: Optional[str] = None,
|
||||
context: dict[str, Any] | None = None,
|
||||
recipient: str | None = None,
|
||||
) -> bool:
|
||||
"""发送严重通知"""
|
||||
notification = Notification(
|
||||
@@ -344,8 +344,8 @@ class NotificationPlugin(BasePlugin):
|
||||
self,
|
||||
provider: str,
|
||||
status: str,
|
||||
error: Optional[str] = None,
|
||||
latency: Optional[float] = None,
|
||||
error: str | None = None,
|
||||
latency: float | None = None,
|
||||
) -> bool:
|
||||
"""发送提供商状态通知"""
|
||||
level = NotificationLevel.INFO
|
||||
@@ -370,7 +370,7 @@ class NotificationPlugin(BasePlugin):
|
||||
)
|
||||
return await self.send(notification)
|
||||
|
||||
async def get_stats(self) -> Dict[str, Any]:
|
||||
async def get_stats(self) -> dict[str, Any]:
|
||||
"""
|
||||
获取统计信息
|
||||
|
||||
@@ -404,7 +404,7 @@ class NotificationPlugin(BasePlugin):
|
||||
|
||||
return base_stats
|
||||
|
||||
async def _get_extra_stats(self) -> Dict[str, Any]:
|
||||
async def _get_extra_stats(self) -> dict[str, Any]:
|
||||
"""
|
||||
获取子类特定的统计信息(子类可选重写)
|
||||
|
||||
|
||||
@@ -4,12 +4,10 @@
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import smtplib
|
||||
from datetime import datetime
|
||||
from email.mime.multipart import MIMEMultipart
|
||||
from email.mime.text import MIMEText
|
||||
from typing import Any, Dict, List, Optional
|
||||
from typing import Any
|
||||
|
||||
aiosmtplib: Any
|
||||
try:
|
||||
@@ -33,7 +31,7 @@ class EmailNotificationPlugin(NotificationPlugin):
|
||||
支持HTML和纯文本邮件
|
||||
"""
|
||||
|
||||
def __init__(self, name: str = "email", config: Optional[Dict[str, Any]] = None):
|
||||
def __init__(self, name: str = "email", config: dict[str, Any] | None = None):
|
||||
super().__init__(name, config or {})
|
||||
|
||||
# SMTP配置
|
||||
@@ -60,9 +58,9 @@ class EmailNotificationPlugin(NotificationPlugin):
|
||||
)
|
||||
|
||||
# 缓冲配置
|
||||
self._buffer: List[Notification] = []
|
||||
self._buffer: list[Notification] = []
|
||||
self._lock = asyncio.Lock()
|
||||
self._flush_task: Optional[asyncio.Task[None]] = None
|
||||
self._flush_task: asyncio.Task[None] | None = None
|
||||
|
||||
# 验证配置
|
||||
config_errors = []
|
||||
@@ -116,7 +114,7 @@ class EmailNotificationPlugin(NotificationPlugin):
|
||||
logger.warning("Email 插件刷新任务等待事件循环创建")
|
||||
pass
|
||||
|
||||
def _format_html_email(self, notifications: List[Notification]) -> str:
|
||||
def _format_html_email(self, notifications: list[Notification]) -> str:
|
||||
"""格式化HTML邮件"""
|
||||
# 颜色映射
|
||||
color_map = {
|
||||
@@ -172,7 +170,7 @@ class EmailNotificationPlugin(NotificationPlugin):
|
||||
|
||||
return html
|
||||
|
||||
def _format_text_email(self, notifications: List[Notification]) -> str:
|
||||
def _format_text_email(self, notifications: list[Notification]) -> str:
|
||||
"""格式化纯文本邮件"""
|
||||
lines = ["Notifications from Aether", "=" * 50, ""]
|
||||
|
||||
@@ -303,7 +301,7 @@ class EmailNotificationPlugin(NotificationPlugin):
|
||||
|
||||
return True
|
||||
|
||||
async def _do_send_batch(self, notifications: List[Notification]) -> Dict[str, int]:
|
||||
async def _do_send_batch(self, notifications: list[Notification]) -> dict[str, int]:
|
||||
"""实际批量发送通知"""
|
||||
if not notifications:
|
||||
return {"total": 0, "sent": 0, "failed": 0}
|
||||
@@ -348,7 +346,7 @@ class EmailNotificationPlugin(NotificationPlugin):
|
||||
async with self._lock:
|
||||
return await self._flush_buffer()
|
||||
|
||||
async def _get_extra_stats(self) -> Dict[str, Any]:
|
||||
async def _get_extra_stats(self) -> dict[str, Any]:
|
||||
"""获取 Email 特定的统计信息"""
|
||||
return {
|
||||
"type": "email",
|
||||
|
||||
@@ -3,13 +3,18 @@ Webhook通知插件
|
||||
通过HTTP Webhook发送通知
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import hmac
|
||||
import json
|
||||
import time
|
||||
from datetime import datetime
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import aiohttp
|
||||
|
||||
try:
|
||||
import aiohttp
|
||||
@@ -17,7 +22,6 @@ try:
|
||||
AIOHTTP_AVAILABLE = True
|
||||
except ImportError:
|
||||
AIOHTTP_AVAILABLE = False
|
||||
aiohttp = None
|
||||
|
||||
from src.core.logger import logger
|
||||
|
||||
@@ -30,7 +34,7 @@ class WebhookNotificationPlugin(NotificationPlugin):
|
||||
支持多种Webhook格式(Slack, Discord, 通用)
|
||||
"""
|
||||
|
||||
def __init__(self, name: str = "webhook", config: Dict[str, Any] = None):
|
||||
def __init__(self, name: str = "webhook", config: dict[str, Any] = None):
|
||||
super().__init__(name, config)
|
||||
|
||||
if not AIOHTTP_AVAILABLE:
|
||||
@@ -48,9 +52,9 @@ class WebhookNotificationPlugin(NotificationPlugin):
|
||||
self.headers = config.get("headers", {}) if config else {}
|
||||
|
||||
# 缓冲配置
|
||||
self._buffer: List[Notification] = []
|
||||
self._buffer: list[Notification] = []
|
||||
self._lock = asyncio.Lock()
|
||||
self._session: Optional[aiohttp.ClientSession] = None
|
||||
self._session: aiohttp.ClientSession | None = None
|
||||
self._flush_task = None
|
||||
|
||||
if not self.webhook_url:
|
||||
@@ -89,7 +93,7 @@ class WebhookNotificationPlugin(NotificationPlugin):
|
||||
|
||||
return signature
|
||||
|
||||
def _format_for_slack(self, notification: Notification) -> Dict[str, Any]:
|
||||
def _format_for_slack(self, notification: Notification) -> dict[str, Any]:
|
||||
"""格式化为Slack消息"""
|
||||
# Slack颜色映射
|
||||
color_map = {
|
||||
@@ -120,7 +124,7 @@ class WebhookNotificationPlugin(NotificationPlugin):
|
||||
],
|
||||
}
|
||||
|
||||
def _format_for_discord(self, notification: Notification) -> Dict[str, Any]:
|
||||
def _format_for_discord(self, notification: Notification) -> dict[str, Any]:
|
||||
"""格式化为Discord消息"""
|
||||
# Discord颜色映射
|
||||
color_map = {
|
||||
@@ -150,7 +154,7 @@ class WebhookNotificationPlugin(NotificationPlugin):
|
||||
|
||||
return {"embeds": embeds}
|
||||
|
||||
def _format_for_teams(self, notification: Notification) -> Dict[str, Any]:
|
||||
def _format_for_teams(self, notification: Notification) -> dict[str, Any]:
|
||||
"""格式化为Microsoft Teams消息"""
|
||||
# Teams颜色映射
|
||||
color_map = {
|
||||
@@ -176,7 +180,7 @@ class WebhookNotificationPlugin(NotificationPlugin):
|
||||
"summary": notification.title,
|
||||
}
|
||||
|
||||
def _format_payload(self, notification: Notification) -> Dict[str, Any]:
|
||||
def _format_payload(self, notification: Notification) -> dict[str, Any]:
|
||||
"""根据Webhook类型格式化负载"""
|
||||
if self.webhook_type == "slack":
|
||||
return self._format_for_slack(notification)
|
||||
@@ -208,7 +212,7 @@ class WebhookNotificationPlugin(NotificationPlugin):
|
||||
|
||||
return True
|
||||
|
||||
async def _do_send_batch(self, notifications: List[Notification]) -> Dict[str, Any]:
|
||||
async def _do_send_batch(self, notifications: list[Notification]) -> dict[str, Any]:
|
||||
"""实际批量发送通知"""
|
||||
success_count = 0
|
||||
failed_count = 0
|
||||
@@ -272,7 +276,7 @@ class WebhookNotificationPlugin(NotificationPlugin):
|
||||
async with self._lock:
|
||||
return await self._flush_buffer()
|
||||
|
||||
async def _get_extra_stats(self) -> Dict[str, Any]:
|
||||
async def _get_extra_stats(self) -> dict[str, Any]:
|
||||
"""获取 Webhook 特定的统计信息"""
|
||||
return {
|
||||
"type": "webhook",
|
||||
|
||||
@@ -3,12 +3,12 @@
|
||||
定义速率限制策略的接口
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from abc import abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from typing import Any, Dict, List, Optional
|
||||
from typing import Any
|
||||
|
||||
from ..common import BasePlugin, HealthStatus, PluginMetadata
|
||||
from ..common import BasePlugin
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -19,10 +19,10 @@ class RateLimitResult:
|
||||
|
||||
allowed: bool
|
||||
remaining: int
|
||||
reset_at: Optional[datetime] = None
|
||||
retry_after: Optional[int] = None
|
||||
message: Optional[str] = None
|
||||
headers: Optional[Dict[str, str]] = None
|
||||
reset_at: datetime | None = None
|
||||
retry_after: int | None = None
|
||||
message: str | None = None
|
||||
headers: dict[str, str] | None = None
|
||||
|
||||
def __post_init__(self):
|
||||
if self.headers is None:
|
||||
@@ -49,9 +49,9 @@ class RateLimitStrategy(BasePlugin):
|
||||
author: str = "Unknown",
|
||||
description: str = "",
|
||||
api_version: str = "1.0",
|
||||
dependencies: List[str] = None,
|
||||
provides: List[str] = None,
|
||||
config: Dict[str, Any] = None,
|
||||
dependencies: list[str] = None,
|
||||
provides: list[str] = None,
|
||||
config: dict[str, Any] = None,
|
||||
):
|
||||
"""
|
||||
初始化速率限制策略
|
||||
@@ -119,7 +119,7 @@ class RateLimitStrategy(BasePlugin):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def get_stats(self, key: str) -> Dict[str, Any]:
|
||||
async def get_stats(self, key: str) -> dict[str, Any]:
|
||||
"""
|
||||
获取统计信息
|
||||
|
||||
|
||||
@@ -22,7 +22,7 @@ import asyncio
|
||||
import time
|
||||
from collections import deque
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Deque, Dict
|
||||
from typing import Any
|
||||
|
||||
from src.core.logger import logger
|
||||
from .base import RateLimitResult, RateLimitStrategy
|
||||
@@ -42,7 +42,7 @@ class SlidingWindow:
|
||||
"""
|
||||
self.window_size = window_size
|
||||
self.max_requests = max_requests
|
||||
self.requests: Deque[float] = deque()
|
||||
self.requests: deque[float] = deque()
|
||||
self.last_access_time: float = time.time()
|
||||
|
||||
def _cleanup(self):
|
||||
@@ -121,7 +121,7 @@ class SlidingWindowStrategy(RateLimitStrategy):
|
||||
|
||||
def __init__(self):
|
||||
super().__init__("sliding_window")
|
||||
self.windows: Dict[str, SlidingWindow] = {}
|
||||
self.windows: dict[str, SlidingWindow] = {}
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
# 默认配置
|
||||
@@ -300,7 +300,7 @@ class SlidingWindowStrategy(RateLimitStrategy):
|
||||
|
||||
logger.info(f"滑动窗口已重置")
|
||||
|
||||
async def get_stats(self, key: str) -> Dict[str, Any]:
|
||||
async def get_stats(self, key: str) -> dict[str, Any]:
|
||||
"""
|
||||
获取统计信息
|
||||
|
||||
@@ -324,7 +324,7 @@ class SlidingWindowStrategy(RateLimitStrategy):
|
||||
"reset_at": window.get_reset_time().isoformat(),
|
||||
}
|
||||
|
||||
def configure(self, config: Dict[str, Any]):
|
||||
def configure(self, config: dict[str, Any]):
|
||||
"""
|
||||
配置策略
|
||||
|
||||
@@ -344,7 +344,7 @@ class SlidingWindowStrategy(RateLimitStrategy):
|
||||
self.window_expiry = config.get("window_expiry", self.window_expiry)
|
||||
self._cleanup_interval = config.get("cleanup_interval", self._cleanup_interval)
|
||||
|
||||
def get_memory_stats(self) -> Dict[str, Any]:
|
||||
def get_memory_stats(self) -> dict[str, Any]:
|
||||
"""
|
||||
获取内存使用统计信息
|
||||
|
||||
|
||||
@@ -4,7 +4,7 @@ import asyncio
|
||||
import os
|
||||
import time
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Any, Dict, Optional, Tuple
|
||||
from typing import Any
|
||||
|
||||
from ...clients.redis_client import get_redis_client_sync
|
||||
from src.core.logger import logger
|
||||
@@ -82,7 +82,7 @@ class TokenBucketStrategy(RateLimitStrategy):
|
||||
|
||||
def __init__(self):
|
||||
super().__init__("token_bucket")
|
||||
self.buckets: Dict[str, TokenBucket] = {}
|
||||
self.buckets: dict[str, TokenBucket] = {}
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
# 默认配置
|
||||
@@ -90,11 +90,11 @@ class TokenBucketStrategy(RateLimitStrategy):
|
||||
self.default_refill_rate = 10 # 默认每秒补充10个令牌
|
||||
|
||||
# 可选的 Redis 后端
|
||||
self._redis_backend: Optional[RedisTokenBucketBackend] = None
|
||||
self._redis_backend: RedisTokenBucketBackend | None = None
|
||||
self._redis_checked = False
|
||||
self._backend_mode = os.getenv("RATE_LIMIT_BACKEND", "auto").lower()
|
||||
|
||||
def _get_bucket(self, key: str, rate_limit: Optional[int] = None) -> TokenBucket:
|
||||
def _get_bucket(self, key: str, rate_limit: int | None = None) -> TokenBucket:
|
||||
"""
|
||||
获取或创建令牌桶
|
||||
|
||||
@@ -248,7 +248,7 @@ class TokenBucketStrategy(RateLimitStrategy):
|
||||
|
||||
logger.info(f"令牌桶已重置")
|
||||
|
||||
async def get_stats(self, key: str) -> Dict[str, Any]:
|
||||
async def get_stats(self, key: str) -> dict[str, Any]:
|
||||
"""
|
||||
获取统计信息
|
||||
|
||||
@@ -278,7 +278,7 @@ class TokenBucketStrategy(RateLimitStrategy):
|
||||
"reset_at": bucket.get_reset_time().isoformat(),
|
||||
}
|
||||
|
||||
def configure(self, config: Dict[str, Any]):
|
||||
def configure(self, config: dict[str, Any]):
|
||||
"""
|
||||
配置策略
|
||||
|
||||
@@ -292,7 +292,7 @@ class TokenBucketStrategy(RateLimitStrategy):
|
||||
self.default_capacity = config.get("default_capacity", self.default_capacity)
|
||||
self.default_refill_rate = config.get("default_refill_rate", self.default_refill_rate)
|
||||
|
||||
def _resolve_capacity(self, key: str, rate_limit: Optional[int] = None) -> int:
|
||||
def _resolve_capacity(self, key: str, rate_limit: int | None = None) -> int:
|
||||
if rate_limit is not None:
|
||||
return rate_limit
|
||||
if key.startswith("api_key:"):
|
||||
@@ -301,7 +301,7 @@ class TokenBucketStrategy(RateLimitStrategy):
|
||||
return self.config.get("user_capacity", self.default_capacity * 2)
|
||||
return self.default_capacity
|
||||
|
||||
def _resolve_refill_rate(self, key: str, rate_limit: Optional[int] = None) -> float:
|
||||
def _resolve_refill_rate(self, key: str, rate_limit: int | None = None) -> float:
|
||||
if rate_limit is not None:
|
||||
return rate_limit / 60.0
|
||||
if key.startswith("api_key:"):
|
||||
@@ -404,7 +404,7 @@ class RedisTokenBucketBackend:
|
||||
capacity: int,
|
||||
refill_rate: float,
|
||||
amount: int,
|
||||
) -> Tuple[bool, int]:
|
||||
) -> tuple[bool, int]:
|
||||
result = await self._consume_script(
|
||||
keys=[self._redis_key(key)],
|
||||
args=[time.time(), capacity, refill_rate, amount],
|
||||
@@ -416,7 +416,7 @@ class RedisTokenBucketBackend:
|
||||
async def reset(self, key: str):
|
||||
await self.redis.delete(self._redis_key(key))
|
||||
|
||||
async def get_stats(self, key: str, capacity: int, refill_rate: float) -> Dict[str, Any]:
|
||||
async def get_stats(self, key: str, capacity: int, refill_rate: float) -> dict[str, Any]:
|
||||
data = await self.redis.hmget(self._redis_key(key), "tokens", "timestamp")
|
||||
tokens = data[0]
|
||||
timestamp = data[1]
|
||||
|
||||
@@ -3,9 +3,10 @@ Token计数插件基类
|
||||
定义Token计数的接口
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
from abc import abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
from typing import Any
|
||||
|
||||
from src.plugins.common import BasePlugin
|
||||
|
||||
@@ -21,7 +22,7 @@ class TokenUsage:
|
||||
cache_write_tokens: int = 0 # Claude缓存写入
|
||||
reasoning_tokens: int = 0 # OpenAI o1推理令牌
|
||||
|
||||
def __add__(self, other: "TokenUsage") -> "TokenUsage":
|
||||
def __add__(self, other: TokenUsage) -> TokenUsage:
|
||||
"""令牌使用相加"""
|
||||
return TokenUsage(
|
||||
input_tokens=self.input_tokens + other.input_tokens,
|
||||
@@ -32,7 +33,7 @@ class TokenUsage:
|
||||
reasoning_tokens=self.reasoning_tokens + other.reasoning_tokens,
|
||||
)
|
||||
|
||||
def to_dict(self) -> Dict[str, int]:
|
||||
def to_dict(self) -> dict[str, int]:
|
||||
"""转换为字典"""
|
||||
return {
|
||||
"input_tokens": self.input_tokens,
|
||||
@@ -50,7 +51,7 @@ class TokenCounterPlugin(BasePlugin):
|
||||
支持不同模型的Token计数
|
||||
"""
|
||||
|
||||
def __init__(self, name: str = "token_counter", config: Dict[str, Any] = None):
|
||||
def __init__(self, name: str = "token_counter", config: dict[str, Any] = None):
|
||||
# 调用父类初始化,设置metadata
|
||||
super().__init__(
|
||||
name=name, config=config, description="Token Counter Plugin", version="1.0.0"
|
||||
@@ -65,25 +66,25 @@ class TokenCounterPlugin(BasePlugin):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def count_tokens(self, text: str, model: Optional[str] = None) -> int:
|
||||
async def count_tokens(self, text: str, model: str | None = None) -> int:
|
||||
"""计算文本的Token数量"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def count_messages(
|
||||
self, messages: List[Dict[str, Any]], model: Optional[str] = None
|
||||
self, messages: list[dict[str, Any]], model: str | None = None
|
||||
) -> int:
|
||||
"""计算消息列表的Token数量"""
|
||||
pass
|
||||
|
||||
async def count_request(self, request: Dict[str, Any], model: Optional[str] = None) -> int:
|
||||
async def count_request(self, request: dict[str, Any], model: str | None = None) -> int:
|
||||
"""计算请求的Token数量"""
|
||||
model = model or request.get("model") or self.default_model
|
||||
messages = request.get("messages", [])
|
||||
return await self.count_messages(messages, model)
|
||||
|
||||
async def count_response(
|
||||
self, response: Dict[str, Any], model: Optional[str] = None
|
||||
self, response: dict[str, Any], model: str | None = None
|
||||
) -> TokenUsage:
|
||||
"""从响应中提取Token使用情况"""
|
||||
usage = response.get("usage", {})
|
||||
@@ -112,8 +113,8 @@ class TokenCounterPlugin(BasePlugin):
|
||||
return TokenUsage()
|
||||
|
||||
async def estimate_cost(
|
||||
self, usage: TokenUsage, model: str, provider: Optional[str] = None
|
||||
) -> Dict[str, float]:
|
||||
self, usage: TokenUsage, model: str, provider: str | None = None
|
||||
) -> dict[str, float]:
|
||||
"""估算使用成本"""
|
||||
# 默认价格表(每1M tokens的价格)
|
||||
pricing = self.config.get("pricing", {})
|
||||
@@ -156,11 +157,11 @@ class TokenCounterPlugin(BasePlugin):
|
||||
}
|
||||
|
||||
@abstractmethod
|
||||
async def get_model_info(self, model: str) -> Dict[str, Any]:
|
||||
async def get_model_info(self, model: str) -> dict[str, Any]:
|
||||
"""获取模型信息"""
|
||||
pass
|
||||
|
||||
async def get_stats(self) -> Dict[str, Any]:
|
||||
async def get_stats(self) -> dict[str, Any]:
|
||||
"""获取统计信息"""
|
||||
return {
|
||||
"type": self.name,
|
||||
|
||||
@@ -5,9 +5,9 @@ Claude Token计数插件
|
||||
|
||||
import json
|
||||
import re
|
||||
from typing import Any, Dict, List, Optional
|
||||
from typing import Any
|
||||
|
||||
from .base import TokenCounterPlugin, TokenUsage
|
||||
from .base import TokenCounterPlugin
|
||||
|
||||
|
||||
class ClaudeTokenCounterPlugin(TokenCounterPlugin):
|
||||
@@ -61,7 +61,7 @@ class ClaudeTokenCounterPlugin(TokenCounterPlugin):
|
||||
},
|
||||
}
|
||||
|
||||
def __init__(self, name: str = "claude", config: Dict[str, Any] = None):
|
||||
def __init__(self, name: str = "claude", config: dict[str, Any] = None):
|
||||
super().__init__(name, config)
|
||||
|
||||
# 价格表(每1M tokens的价格 USD)
|
||||
@@ -153,7 +153,7 @@ class ClaudeTokenCounterPlugin(TokenCounterPlugin):
|
||||
# 取两者的平均
|
||||
return (token_by_words + token_by_chars) // 2
|
||||
|
||||
async def count_tokens(self, text: str, model: Optional[str] = None) -> int:
|
||||
async def count_tokens(self, text: str, model: str | None = None) -> int:
|
||||
"""计算文本的Token数量"""
|
||||
if not self.enabled:
|
||||
return 0
|
||||
@@ -162,7 +162,7 @@ class ClaudeTokenCounterPlugin(TokenCounterPlugin):
|
||||
return self._estimate_tokens_from_text(text, model)
|
||||
|
||||
async def count_messages(
|
||||
self, messages: List[Dict[str, Any]], model: Optional[str] = None
|
||||
self, messages: list[dict[str, Any]], model: str | None = None
|
||||
) -> int:
|
||||
"""计算消息列表的Token数量"""
|
||||
if not self.enabled:
|
||||
@@ -213,7 +213,7 @@ class ClaudeTokenCounterPlugin(TokenCounterPlugin):
|
||||
|
||||
return total_tokens
|
||||
|
||||
async def count_request(self, request: Dict[str, Any], model: Optional[str] = None) -> int:
|
||||
async def count_request(self, request: dict[str, Any], model: str | None = None) -> int:
|
||||
"""计算请求的Token数量"""
|
||||
model = model or request.get("model") or self.default_model
|
||||
messages = request.get("messages", [])
|
||||
@@ -227,7 +227,7 @@ class ClaudeTokenCounterPlugin(TokenCounterPlugin):
|
||||
|
||||
return total
|
||||
|
||||
async def get_model_info(self, model: str) -> Dict[str, Any]:
|
||||
async def get_model_info(self, model: str) -> dict[str, Any]:
|
||||
"""获取模型信息"""
|
||||
info = {"model": model, "supported": self.supports_model(model)}
|
||||
|
||||
@@ -261,7 +261,7 @@ class ClaudeTokenCounterPlugin(TokenCounterPlugin):
|
||||
|
||||
return info
|
||||
|
||||
async def get_stats(self) -> Dict[str, Any]:
|
||||
async def get_stats(self) -> dict[str, Any]:
|
||||
"""获取统计信息"""
|
||||
stats = await super().get_stats()
|
||||
stats.update(
|
||||
|
||||
@@ -3,12 +3,11 @@ Tiktoken Token计数插件
|
||||
支持OpenAI和其他使用tiktoken的模型
|
||||
"""
|
||||
|
||||
import json
|
||||
from typing import Any, Dict, List, Optional
|
||||
from typing import Any
|
||||
|
||||
from src.core.logger import logger
|
||||
|
||||
from .base import TokenCounterPlugin, TokenUsage
|
||||
from .base import TokenCounterPlugin
|
||||
|
||||
# 尝试导入tiktoken
|
||||
try:
|
||||
@@ -57,7 +56,7 @@ class TiktokenCounterPlugin(TokenCounterPlugin):
|
||||
"gpt-4o-mini": 3,
|
||||
}
|
||||
|
||||
def __init__(self, name: str = "tiktoken", config: Dict[str, Any] = None):
|
||||
def __init__(self, name: str = "tiktoken", config: dict[str, Any] = None):
|
||||
super().__init__(name, config)
|
||||
|
||||
if not TIKTOKEN_AVAILABLE:
|
||||
@@ -121,7 +120,7 @@ class TiktokenCounterPlugin(TokenCounterPlugin):
|
||||
openai_models = ["gpt-4", "gpt-3.5", "text-davinci", "text-embedding", "code-davinci", "o1"]
|
||||
return any(model.startswith(prefix) for prefix in openai_models)
|
||||
|
||||
async def count_tokens(self, text: str, model: Optional[str] = None) -> int:
|
||||
async def count_tokens(self, text: str, model: str | None = None) -> int:
|
||||
"""计算文本的Token数量"""
|
||||
if not self.enabled:
|
||||
return 0
|
||||
@@ -138,7 +137,7 @@ class TiktokenCounterPlugin(TokenCounterPlugin):
|
||||
return int(len(text) * 0.75)
|
||||
|
||||
async def count_messages(
|
||||
self, messages: List[Dict[str, Any]], model: Optional[str] = None
|
||||
self, messages: list[dict[str, Any]], model: str | None = None
|
||||
) -> int:
|
||||
"""计算消息列表的Token数量"""
|
||||
if not self.enabled:
|
||||
@@ -203,7 +202,7 @@ class TiktokenCounterPlugin(TokenCounterPlugin):
|
||||
|
||||
return total_tokens
|
||||
|
||||
async def get_model_info(self, model: str) -> Dict[str, Any]:
|
||||
async def get_model_info(self, model: str) -> dict[str, Any]:
|
||||
"""获取模型信息"""
|
||||
info = {"model": model, "supported": self.supports_model(model)}
|
||||
|
||||
@@ -260,7 +259,7 @@ class TiktokenCounterPlugin(TokenCounterPlugin):
|
||||
# 默认值
|
||||
return 4096
|
||||
|
||||
async def get_stats(self) -> Dict[str, Any]:
|
||||
async def get_stats(self) -> dict[str, Any]:
|
||||
"""获取统计信息"""
|
||||
stats = await super().get_stats()
|
||||
stats.update(
|
||||
|
||||
Reference in New Issue
Block a user