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

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 StrEnum 

17 

18 

19class TenantStatus(StrEnum): 

20 ACTIVE = "active" 

21 SUSPENDED = "suspended" 

22 DELETED = "deleted" 

23 

24 

25class TenantTier(StrEnum): 

26 """租户等级。""" 

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 

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

71 default_provider: str = "openai" 

72 allowed_providers: list[str] = field( 

73 default_factory=lambda: ["openai", "deepseek", "anthropic"] 

74 ) 

75 max_iterations: int = 10 

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

77 custom_settings: dict = field(default_factory=dict) 

78 

79 

80@dataclass 

81class TenantUsage: 

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

83 

84 tenant_id: str 

85 date: str # YYYY-MM-DD 

86 api_calls: int = 0 

87 tokens_used: int = 0 

88 current_concurrency: int = 0 

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

90 

91 

92@dataclass 

93class Tenant: 

94 """租户实体。""" 

95 

96 tenant_id: str 

97 name: str 

98 tier: TenantTier 

99 status: TenantStatus = TenantStatus.ACTIVE 

100 config: TenantConfig = field(default_factory=TenantConfig) 

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

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

103 metadata: dict = field(default_factory=dict) 

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

105 custom_quotas: dict = field(default_factory=dict) 

106 

107 

108class TenantManager: 

109 """多租户管理器。 

110 

111 特性： 

112 - 租户 CRUD + 启停 

113 - 等级配额自动分配 

114 - 用量追踪 + 超限拦截 

115 - 租户级配置隔离 

116 - 每日用量自动重置 

117 """ 

118 

119 def __init__(self): 

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

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

122 

123 # ── 租户管理 ── 

124 

125 def create_tenant( 

126 self, 

127 name: str, 

128 tier: TenantTier = TenantTier.FREE, 

129 config: TenantConfig | None = None, 

130 metadata: dict = None, 

131 ) -> Tenant: 

132 """创建租户。""" 

133 import uuid 

134 

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

136 tenant = Tenant( 

137 tenant_id=tenant_id, 

138 name=name, 

139 tier=tier, 

140 config=config or TenantConfig(), 

141 metadata=metadata or {}, 

142 ) 

143 self._tenants[tenant_id] = tenant 

144 return tenant 

145 

146 def get_tenant(self, tenant_id: str) -> Tenant | None: 

147 return self._tenants.get(tenant_id) 

148 

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

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

151 if status: 

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

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

154 

155 def update_tenant(self, tenant_id: str, **kwargs) -> Tenant | None: 

156 tenant = self._tenants.get(tenant_id) 

157 if not tenant: 

158 return None 

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

160 if hasattr(tenant, k): 

161 setattr(tenant, k, v) 

162 tenant.updated_at = time.time() 

163 return tenant 

164 

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

166 t = self._tenants.get(tenant_id) 

167 if not t: 

168 return False 

169 t.status = TenantStatus.SUSPENDED 

170 t.updated_at = time.time() 

171 return True 

172 

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

174 t = self._tenants.get(tenant_id) 

175 if not t: 

176 return False 

177 t.status = TenantStatus.ACTIVE 

178 t.updated_at = time.time() 

179 return True 

180 

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

182 t = self._tenants.get(tenant_id) 

183 if not t: 

184 return False 

185 t.status = TenantStatus.DELETED 

186 t.updated_at = time.time() 

187 return True 

188 

189 # ── 配额 ── 

190 

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

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

193 tenant = self._tenants.get(tenant_id) 

194 if not tenant: 

195 return {} 

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

197 base.update(tenant.custom_quotas) 

198 return base 

199 

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

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

202 tenant = self._tenants.get(tenant_id) 

203 if not tenant: 

204 return False, "租户不存在" 

205 if tenant.status != TenantStatus.ACTIVE: 

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

207 

208 quotas = self.get_quotas(tenant_id) 

209 limit = quotas.get(resource) 

210 

211 if limit is None: 

212 return True, "" 

213 

214 usage = self._get_usage(tenant_id) 

215 current = getattr(usage, resource, 0) 

216 

217 if current + amount > limit: 

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

219 

220 return True, "" 

221 

222 # ── 用量追踪 ── 

223 

224 def record_usage( 

225 self, tenant_id: str, api_calls: int = 0, tokens: int = 0, concurrency_delta: int = 0 

226 ): 

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

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

229 return 

230 usage = self._get_usage(tenant_id) 

231 usage.api_calls += api_calls 

232 usage.tokens_used += tokens 

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

234 usage.last_updated = time.time() 

235 

236 def get_usage(self, tenant_id: str) -> TenantUsage | None: 

237 return self._get_usage(tenant_id) 

238 

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

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

241 if tenant_id: 

242 self._usage.pop(tenant_id, None) 

243 else: 

244 self._usage.clear() 

245 

246 # ── 统计 ── 

247 

248 def stats(self) -> dict: 

249 total = len(self._tenants) 

250 by_tier = {} 

251 by_status = {} 

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

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

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

255 return { 

256 "total": total, 

257 "by_tier": by_tier, 

258 "by_status": by_status, 

259 } 

260 

261 # ── 内部 ── 

262 

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

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

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

266 if key not in self._usage: 

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

268 return self._usage[key]