Coverage for agentos/enterprise/tenants.py: 42%

1""" 

2AgentOS Enterprise — Multi-Tenant Management. 

3 

4功能： 

5 - 租户创建/启停/删除 

6 - 租户级配额管理（API 调用数、Token 数、并发数） 

7 - 租户隔离（数据/配置/Agent 命名空间） 

8 - 用量追踪与超限拦截 

9 - 租户级自定义配置 

10""" 

11 

12from __future__ import annotations 

13 

14import time 

15from dataclasses import dataclass, field 

16from enum import Enum 

17from typing import Optional 

18 

19 

20class TenantStatus(str, Enum): 

21 ACTIVE = "active" 

22 SUSPENDED = "suspended" 

23 DELETED = "deleted" 

24 

25 

26class TenantTier(str, Enum): 

27 """租户等级。""" 

28 FREE = "free" # 100 调用/天, 1 并发 

29 STARTER = "starter" # 1000 调用/天, 3 并发 

30 PRO = "pro" # 10000 调用/天, 10 并发 

31 ENTERPRISE = "enterprise" # 自定义 

32 

33 

34TIER_QUOTAS = { 

35 TenantTier.FREE: { 

36 "daily_api_calls": 100, 

37 "daily_tokens": 100_000, 

38 "max_concurrency": 1, 

39 "max_agents": 3, 

40 "max_api_keys": 2, 

41 }, 

42 TenantTier.STARTER: { 

43 "daily_api_calls": 1_000, 

44 "daily_tokens": 1_000_000, 

45 "max_concurrency": 3, 

46 "max_agents": 10, 

47 "max_api_keys": 5, 

48 }, 

49 TenantTier.PRO: { 

50 "daily_api_calls": 10_000, 

51 "daily_tokens": 10_000_000, 

52 "max_concurrency": 10, 

53 "max_agents": 50, 

54 "max_api_keys": 20, 

55 }, 

56 TenantTier.ENTERPRISE: { 

57 "daily_api_calls": 1_000_000, 

58 "daily_tokens": 1_000_000_000, 

59 "max_concurrency": 100, 

60 "max_agents": 500, 

61 "max_api_keys": 100, 

62 }, 

63} 

64 

65 

66@dataclass 

67class TenantConfig: 

68 """租户级配置覆盖。""" 

69 default_model: str = "gpt-4o-mini" 

70 default_provider: str = "openai" 

71 allowed_providers: list[str] = field(default_factory=lambda: ["openai", "deepseek", "anthropic"]) 

72 max_iterations: int = 10 

73 guardrail_level: str = "standard" # none / standard / strict 

74 custom_settings: dict = field(default_factory=dict) 

75 

76 

77@dataclass 

78class TenantUsage: 

79 """租户用量统计（当日）。""" 

80 tenant_id: str 

81 date: str # YYYY-MM-DD 

82 api_calls: int = 0 

83 tokens_used: int = 0 

84 current_concurrency: int = 0 

85 last_updated: float = field(default_factory=time.time) 

86 

87 

88@dataclass 

89class Tenant: 

90 """租户实体。""" 

91 tenant_id: str 

92 name: str 

93 tier: TenantTier 

94 status: TenantStatus = TenantStatus.ACTIVE 

95 config: TenantConfig = field(default_factory=TenantConfig) 

96 created_at: float = field(default_factory=time.time) 

97 updated_at: float = field(default_factory=time.time) 

98 metadata: dict = field(default_factory=dict) 

99 # 自定义配额覆盖（仅 Enterprise 级别可用） 

100 custom_quotas: dict = field(default_factory=dict) 

101 

102 

103class TenantManager: 

104 """多租户管理器。 

105 

106 特性： 

107 - 租户 CRUD + 启停 

108 - 等级配额自动分配 

109 - 用量追踪 + 超限拦截 

110 - 租户级配置隔离 

111 - 每日用量自动重置 

112 """ 

113 

114 def __init__(self): 

115 self._tenants: dict[str, Tenant] = {} 

116 self._usage: dict[str, TenantUsage] = {} # tenant_id → usage 

117 

118 # ── 租户管理 ── 

119 

120 def create_tenant( 

121 self, 

122 name: str, 

123 tier: TenantTier = TenantTier.FREE, 

124 config: Optional[TenantConfig] = None, 

125 metadata: dict = None, 

126 ) -> Tenant: 

127 """创建租户。""" 

128 import uuid 

129 tenant_id = f"tn_{uuid.uuid4().hex[:12]}" 

130 tenant = Tenant( 

131 tenant_id=tenant_id, 

132 name=name, 

133 tier=tier, 

134 config=config or TenantConfig(), 

135 metadata=metadata or {}, 

136 ) 

137 self._tenants[tenant_id] = tenant 

138 return tenant 

139 

140 def get_tenant(self, tenant_id: str) -> Optional[Tenant]: 

141 return self._tenants.get(tenant_id) 

142 

143 def list_tenants(self, status: Optional[TenantStatus] = None) -> list[Tenant]: 

144 tenants = list(self._tenants.values()) 

145 if status: 

146 tenants = [t for t in tenants if t.status == status] 

147 return sorted(tenants, key=lambda t: t.created_at) 

148 

149 def update_tenant(self, tenant_id: str, **kwargs) -> Optional[Tenant]: 

150 tenant = self._tenants.get(tenant_id) 

151 if not tenant: 

152 return None 

153 for k, v in kwargs.items(): 

154 if hasattr(tenant, k): 

155 setattr(tenant, k, v) 

156 tenant.updated_at = time.time() 

157 return tenant 

158 

159 def suspend_tenant(self, tenant_id: str) -> bool: 

160 t = self._tenants.get(tenant_id) 

161 if not t: 

162 return False 

163 t.status = TenantStatus.SUSPENDED 

164 t.updated_at = time.time() 

165 return True 

166 

167 def activate_tenant(self, tenant_id: str) -> bool: 

168 t = self._tenants.get(tenant_id) 

169 if not t: 

170 return False 

171 t.status = TenantStatus.ACTIVE 

172 t.updated_at = time.time() 

173 return True 

174 

175 def delete_tenant(self, tenant_id: str) -> bool: 

176 t = self._tenants.get(tenant_id) 

177 if not t: 

178 return False 

179 t.status = TenantStatus.DELETED 

180 t.updated_at = time.time() 

181 return True 

182 

183 # ── 配额 ── 

184 

185 def get_quotas(self, tenant_id: str) -> dict: 

186 """获取租户当前有效配额。""" 

187 tenant = self._tenants.get(tenant_id) 

188 if not tenant: 

189 return {} 

190 base = dict(TIER_QUOTAS.get(tenant.tier, {})) 

191 base.update(tenant.custom_quotas) 

192 return base 

193 

194 def check_quota(self, tenant_id: str, resource: str, amount: int = 1) -> tuple[bool, str]: 

195 """检查配额是否允许此次操作。返回 (允许, 原因)。""" 

196 tenant = self._tenants.get(tenant_id) 

197 if not tenant: 

198 return False, "租户不存在" 

199 if tenant.status != TenantStatus.ACTIVE: 

200 return False, f"租户状态: {tenant.status.value}" 

201 

202 quotas = self.get_quotas(tenant_id) 

203 limit = quotas.get(resource) 

204 

205 if limit is None: 

206 return True, "" 

207 

208 usage = self._get_usage(tenant_id) 

209 current = getattr(usage, resource, 0) 

210 

211 if current + amount > limit: 

212 return False, f"超出配额: {resource} ({current}/{limit})" 

213 

214 return True, "" 

215 

216 # ── 用量追踪 ── 

217 

218 def record_usage(self, tenant_id: str, api_calls: int = 0, tokens: int = 0, concurrency_delta: int = 0): 

219 """记录一次用量。""" 

220 if not self._tenants.get(tenant_id): 

221 return 

222 usage = self._get_usage(tenant_id) 

223 usage.api_calls += api_calls 

224 usage.tokens_used += tokens 

225 usage.current_concurrency = max(0, usage.current_concurrency + concurrency_delta) 

226 usage.last_updated = time.time() 

227 

228 def get_usage(self, tenant_id: str) -> Optional[TenantUsage]: 

229 return self._get_usage(tenant_id) 

230 

231 def reset_daily_usage(self, tenant_id: str = None): 

232 """重置每日用量（定时任务调用）。""" 

233 if tenant_id: 

234 self._usage.pop(tenant_id, None) 

235 else: 

236 self._usage.clear() 

237 

238 # ── 统计 ── 

239 

240 def stats(self) -> dict: 

241 total = len(self._tenants) 

242 by_tier = {} 

243 by_status = {} 

244 for t in self._tenants.values(): 

245 by_tier[t.tier.value] = by_tier.get(t.tier.value, 0) + 1 

246 by_status[t.status.value] = by_status.get(t.status.value, 0) + 1 

247 return { 

248 "total": total, 

249 "by_tier": by_tier, 

250 "by_status": by_status, 

251 } 

252 

253 # ── 内部 ── 

254 

255 def _get_usage(self, tenant_id: str) -> TenantUsage: 

256 today = time.strftime("%Y-%m-%d") 

257 key = f"{tenant_id}:{today}" 

258 if key not in self._usage: 

259 self._usage[key] = TenantUsage(tenant_id=tenant_id, date=today) 

260 return self._usage[key]