File size: 13,145 Bytes
a37e6db | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 | """
Platform Management Tools for Atom
Provides administrative tools for managing tenants, workspaces, members, and teams.
"""
import logging
import uuid
import json
from typing import Any, Dict, List, Optional
from datetime import datetime
from core.database import SessionLocal
from core.models import Tenant, Workspace, User, Team, team_members
from sqlalchemy import or_
logger = logging.getLogger(__name__)
async def get_platform_settings(context: Dict[str, Any] = None) -> Dict[str, Any]:
"""Retrieves platform-wide settings and configuration for the current tenant."""
from core.database import SessionLocal
from core.models import TenantSetting
workspace_id = (context.get("workspace_id") if context else None) or "default"
try:
with SessionLocal() as db:
settings = db.query(TenantSetting).filter(TenantSetting.tenant_id == workspace_id).all()
return {s.setting_key: s.setting_value for s in settings}
except Exception as e:
logger.error(f"Error fetching platform settings: {e}")
return {"error": f"Failed to fetch settings: {str(e)}"}
async def update_platform_setting(key: str, value: str, context: Dict[str, Any] = None) -> str:
"""Updates or creates a platform-wide setting."""
from core.database import SessionLocal
from core.models import TenantSetting
workspace_id = (context.get("workspace_id") if context else None) or "default"
try:
with SessionLocal() as db:
setting = db.query(TenantSetting).filter(
TenantSetting.tenant_id == workspace_id,
TenantSetting.setting_key == key
).first()
if setting:
setting.setting_value = value
else:
setting = TenantSetting(tenant_id=workspace_id, setting_key=key, setting_value=value)
db.add(setting)
db.commit()
return f"Setting '{key}' successfully updated to '{value}'."
except Exception as e:
logger.error(f"Error updating platform setting: {e}")
return f"Error: Failed to update setting: {str(e)}"
async def update_tenant_profile(
name: Optional[str] = None,
billing_email: Optional[str] = None,
logo_url: Optional[str] = None,
primary_color: Optional[str] = None,
budget_limit_usd: Optional[float] = None,
context: Dict[str, Any] = None
) -> str:
"""Updates the tenant's profile information."""
from core.database import SessionLocal
from core.models import Tenant, Workspace
workspace_id = (context.get("workspace_id") if context else None) or "default"
try:
with SessionLocal() as db:
# First get tenant_id from workspace
ws = db.query(Workspace).filter(Workspace.id == workspace_id).first()
# Fallback for single-tenant mode where workspace might not exist yet
tenant_id = ws.tenant_id if ws else "default"
tenant = db.query(Tenant).filter(Tenant.id == tenant_id).first()
# If default tenant doesn't exist, we might be in ultra-lean OS mode
if not tenant and tenant_id == "default":
# Create default tenant on the fly if needed? Or just return error
return f"Error: Default tenant not found. Please initialize the system."
elif not tenant:
return f"Error: Tenant {tenant_id} not found."
updates = []
if name:
tenant.name = name
updates.append("name")
if billing_email:
tenant.billing_email = billing_email
updates.append("billing_email")
if logo_url:
metadata = tenant.metadata_json or {}
metadata["logo_url"] = logo_url
tenant.metadata_json = metadata
updates.append("logo_url")
if primary_color:
metadata = tenant.metadata_json or {}
metadata["primary_color"] = primary_color
tenant.metadata_json = metadata
updates.append("primary_color")
if budget_limit_usd is not None:
tenant.budget_limit_usd = budget_limit_usd
updates.append("budget_limit_usd")
if not updates:
return "No updates provided."
db.commit()
return f"Tenant profile updated successfully. Fields modified: {', '.join(updates)}."
except Exception as e:
logger.error(f"Error updating tenant profile: {e}")
return f"Error: Failed to update tenant profile: {str(e)}"
async def set_byok_api_key(
provider: str,
api_key: str,
context: Dict[str, Any] = None
) -> str:
"""Sets or updates a Bring-Your-Own-Key (BYOK) API key for a provider."""
from core.database import SessionLocal
from core.byok_endpoints import BYOKManager
tenant_id = context.get("workspace_id") if context else None
if not tenant_id:
return "Error: Could not resolve tenant/workspace ID from context."
db = SessionLocal()
try:
manager = BYOKManager(db)
await manager.set_api_key(tenant_id, provider, api_key)
db.commit()
return f"Successfully set API key for {provider} in tenant {tenant_id}."
except Exception as e:
db.rollback()
logger.error(f"Error setting BYOK API key: {e}")
return f"Error setting BYOK API key: {str(e)}"
finally:
db.close()
async def list_tenant_members(context: Dict[str, Any] = None) -> str:
"""Lists all members of the current tenant."""
from core.database import SessionLocal
from core.models import User, Workspace
workspace_id = context.get("workspace_id") if context else None
if not workspace_id:
return "Error: Could not resolve workspace ID from context."
db = SessionLocal()
try:
# First get the tenant_id from the workspace
ws = db.query(Workspace).filter(Workspace.id == workspace_id).first()
if not ws:
return f"Error: Workspace {workspace_id} not found."
tenant_id = ws.tenant_id
# Get all users belonging to this tenant
# In multi-tenant architecture, users usually have a tenant_id
members = db.query(User).filter(User.tenant_id == tenant_id).all()
if not members:
return "No members found for this tenant."
result = [f"Members for Tenant {tenant_id}:"]
for m in members:
result.append(f"- {m.full_name or m.email} (ID: {m.id}, Role: {getattr(m, 'role', 'N/A')}, Status: {getattr(m, 'status', 'N/A')})")
return "\n".join(result)
except Exception as e:
logger.error(f"Error listing tenant members: {e}")
return f"Error listing tenant members: {str(e)}"
finally:
db.close()
async def manage_tenant_member(
user_id: str,
action: str, # "update_role", "deactivate", "reactivate"
role: Optional[str] = None,
context: Dict[str, Any] = None
) -> str:
"""Updates a member's role or status within a tenant."""
from core.database import SessionLocal
from core.models import User
db = SessionLocal()
try:
user = db.query(User).filter(User.id == user_id).first()
if not user:
return f"Error: User {user_id} not found."
if action == "update_role":
if not role:
return "Error: role is required for update_role action."
user.role = role
msg = f"User {user_id} role updated to {role}."
elif action == "deactivate":
user.is_active = False
msg = f"User {user_id} deactivated."
elif action == "reactivate":
user.is_active = True
msg = f"User {user_id} reactivated."
else:
return f"Error: Unknown action {action}."
db.commit()
return msg
except Exception as e:
db.rollback()
logger.error(f"Error managing tenant member: {e}")
return f"Error managing tenant member: {str(e)}"
finally:
db.close()
async def manage_workspace(
name: str,
action: str = "create", # "create", "update"
workspace_id: Optional[str] = None,
description: Optional[str] = None,
is_startup: bool = False,
context: Dict[str, Any] = None
) -> str:
"""Creates or updates a workspace within a tenant."""
from core.database import SessionLocal
from core.models import Workspace, Tenant
# Resolve tenant_id
actual_tenant_id = context.get("tenant_id") if context else None
if not actual_tenant_id:
workspace_id_ctx = context.get("workspace_id") if context else None
if workspace_id_ctx:
# Try to lookup tenant from workspace
with SessionLocal() as db_temp:
ws_temp = db_temp.query(Workspace).filter(Workspace.id == workspace_id_ctx).first()
if ws_temp:
actual_tenant_id = ws_temp.tenant_id
if not actual_tenant_id:
return "Error: Could not resolve tenant ID from context."
db = SessionLocal()
try:
if action == "create":
new_ws = Workspace(
name=name,
description=description,
is_startup=is_startup,
tenant_id=actual_tenant_id
)
db.add(new_ws)
db.commit()
return f"Workspace '{name}' created successfully with ID {new_ws.id}."
elif action == "update":
if not workspace_id:
return "Error: workspace_id is required for update action."
ws = db.query(Workspace).filter(Workspace.id == workspace_id).first()
if not ws:
return f"Error: Workspace {workspace_id} not found."
ws.name = name
if description is not None:
ws.description = description
ws.is_startup = is_startup
db.commit()
return f"Workspace {workspace_id} updated successfully."
else:
return f"Error: Unknown action {action}."
except Exception as e:
db.rollback()
logger.error(f"Error managing workspace: {e}")
return f"Error managing workspace: {str(e)}"
finally:
db.close()
async def manage_team(
name: str,
action: str = "create", # "create", "update"
team_id: Optional[str] = None,
workspace_id: Optional[str] = None,
add_members: Optional[List[str]] = None, # List of user emails or IDs
context: Dict[str, Any] = None
) -> str:
"""Creates or updates a Team and its members."""
from core.database import SessionLocal
from core.models import Team, User, team_members, Workspace, Tenant
tenant_id = context.get("workspace_id") if context else None
if not tenant_id:
return "Error: Could not resolve tenant/workspace context."
db = SessionLocal()
try:
# Resolve actual workspace_id if not provided
if not workspace_id:
workspace_id = context.get("workspace_id")
if action == "create":
team = Team(
name=name,
workspace_id=workspace_id
)
db.add(team)
db.flush() # Get team.id
team_id = team.id
msg = f"Team '{name}' created successfully with ID {team_id}."
elif action == "update":
if not team_id:
return "Error: team_id is required for update action."
team = db.query(Team).filter(Team.id == team_id).first()
if not team:
return f"Error: Team {team_id} not found."
team.name = name
msg = f"Team {team_id} updated successfully."
else:
return f"Error: Unknown action {action}."
# Add members if provided
if add_members and team_id:
added = 0
for identifier in add_members:
user = db.query(User).filter(or_(User.id == identifier, User.email == identifier)).first()
if user:
# Check if already in team
existing = db.query(team_members).filter(
team_members.c.user_id == user.id,
team_members.c.team_id == team_id
).first()
if not existing:
db.execute(team_members.insert().values(user_id=user.id, team_id=team_id))
added += 1
msg += f" Added {added} members."
db.commit()
return msg
except Exception as e:
db.rollback()
logger.error(f"Error managing team: {e}")
return f"Error managing team: {str(e)}"
finally:
db.close()
|