Spaces:
Sleeping
Sleeping
| """ | |
| 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() | |