Merge pull request #63 from amazeeio/periodic-job

Periodic job for team management

Author: PhiRho

app/api/private_ai_keys.py

@@ -118,8 +118,9 @@ async def create_vector_db(
     if settings.ENABLE_LIMITS:
         if not team_id: # if the team_id is not set we have already validated the owner_id
             user = db.query(DBUser).filter(DBUser.id == owner_id).first()
-            team_id = user.team_id or FAKE_ID
-        check_vector_db_limits(db, team_id)
+            team_id = user.team_id  # Remove the FAKE_ID fallback
+        if team_id:  # Only check limits if we have a valid team_id
+            check_vector_db_limits(db, team_id)
 
     try:
         # Create new postgres database
@@ -332,7 +333,7 @@ async def create_llm_token(
             litellm_token=litellm_token,
             litellm_api_url=region.litellm_api_url,
             owner_id=owner_id,
-            team_id=team_id,
+            team_id=None if team_id is None else team_id,
             name=private_ai_key.name,
             region_id = private_ai_key.region_id
         )

app/core/config.py

@@ -34,6 +34,7 @@ class Settings(BaseSettings):
 
     model_config = ConfigDict(env_file=".env")
     main_route: str = os.getenv("LAGOON_ROUTE", "http://localhost:8800")
+    frontend_route: str = os.getenv("FRONTEND_ROUTE", "http://localhost:3000")
 
     def model_post_init(self, values):
         # Add Lagoon routes to CORS origins if available

app/core/resource_limits.py

@@ -15,7 +15,7 @@
 DEFAULT_SERVICE_KEYS = 1
 DEFAULT_VECTOR_DB_COUNT = 1
 DEFAULT_KEY_DURATION = 30
-DEFAULT_MAX_SPEND = 20.0
+DEFAULT_MAX_SPEND = 27.0
 DEFAULT_RPM_PER_KEY = 500
 
 def check_team_user_limit(db: Session, team_id: int) -> None:

app/core/worker.py

@@ -1,7 +1,8 @@
-from datetime import datetime, UTC
+from datetime import datetime, UTC, timedelta
 from sqlalchemy.orm import Session
-from app.db.models import DBTeam, DBProduct, DBTeamProduct, DBPrivateAIKey, DBUser
+from app.db.models import DBTeam, DBProduct, DBTeamProduct, DBPrivateAIKey, DBUser, DBRegion
 from app.services.litellm import LiteLLMService
+from app.services.ses import SESService
 import logging
 from collections import defaultdict
 from app.core.resource_limits import get_token_restrictions
@@ -15,9 +16,79 @@
     INVOICE_FAILURE_EVENTS,
     INVOICE_SUCCESS_EVENTS
 )
+from prometheus_client import Gauge, Counter
+from typing import Dict, List
+from app.core.security import create_access_token
+from app.core.config import settings
+from urllib.parse import urljoin
 
 logger = logging.getLogger(__name__)
 
+FIRST_EMAIL_DAYS_LEFT = 7
+SECOND_EMAIL_DAYS_LEFT = 5
+TRIAL_OVER_DAYS = 30
+
+# Prometheus metrics
+team_freshness_days = Gauge(
+    "team_freshness_days",
+    "Freshness of teams in days (since creation for teams without products, since last payment for teams with products)",
+    ["team_id", "team_name"]
+)
+
+team_expired_metric = Counter(
+    "team_expired_total",
+    "Total number of teams that have expired without products",
+    ["team_id", "team_name"]
+)
+
+key_spend_percentage = Gauge(
+    "key_spend_percentage",
+    "Percentage of budget used for each key",
+    ["team_id", "team_name", "key_alias"]
+)
+
+team_total_spend = Gauge(
+    "team_total_spend",
+    "Total spend across all keys in a team for the current budget period",
+    ["team_id", "team_name"]
+)
+
+# Track active team labels to zero out metrics for inactive teams
+active_team_labels = set()
+
+def get_team_keys_by_region(db: Session, team_id: int) -> Dict[DBRegion, List[DBPrivateAIKey]]:
+    """
+    Get all keys for a team grouped by region.
+
+    Args:
+        db: Database session
+        team_id: ID of the team to get keys for
+
+    Returns:
+        Dictionary mapping regions to lists of keys
+    """
+    # Get all keys for the team with their regions
+    team_users = db.query(DBUser).filter(DBUser.team_id == team_id).all()
+    team_user_ids = [user.id for user in team_users]
+    # Return keys owned by users in the team OR owned by the team
+    team_keys = db.query(DBPrivateAIKey).filter(
+        (DBPrivateAIKey.owner_id.in_(team_user_ids)) |
+        (DBPrivateAIKey.team_id == team_id)
+    ).all()
+
+    # Group keys by region
+    keys_by_region = defaultdict(list)
+    for key in team_keys:
+        if not key.litellm_token:
+            logger.warning(f"Key {key.id} has no LiteLLM token, skipping")
+            continue
+        if not key.region:
+            logger.warning(f"Key {key.id} has no region, skipping")
+            continue
+        keys_by_region[key.region].append(key)
+
+    return keys_by_region
+
 async def handle_stripe_event_background(event, db: Session):
     """
     Background task to handle Stripe webhook events.
@@ -60,7 +131,6 @@ async def handle_stripe_event_background(event, db: Session):
     except Exception as e:
         logger.error(f"Error in background event handler: {str(e)}")
 
-
 async def apply_product_for_team(db: Session, customer_id: str, product_id: str, start_date: datetime):
     """
     Apply a product to a team and update their last payment date.
@@ -107,25 +177,8 @@ async def apply_product_for_team(db: Session, customer_id: str, product_id: str,
 
         days_left_in_period, max_max_spend, max_rpm_limit = get_token_restrictions(db, team.id)
 
-        # Get all keys for the team with their regions
-        team_users = db.query(DBUser).filter(DBUser.team_id == team.id).all()
-        team_user_ids = [user.id for user in team_users]
-        # Return keys owned by users in the team OR owned by the team
-        team_keys = db.query(DBPrivateAIKey).filter(
-            (DBPrivateAIKey.owner_id.in_(team_user_ids)) |
-            (DBPrivateAIKey.team_id == team.id)
-        ).all()
-
-        # Group keys by region
-        keys_by_region = defaultdict(list)
-        for key in team_keys:
-            if not key.litellm_token:
-                logger.warning(f"Key {key.id} has no LiteLLM token, skipping")
-                continue
-            if not key.region:
-                logger.warning(f"Key {key.id} has no region, skipping")
-                continue
-            keys_by_region[key.region].append(key)
+        # Get all keys for the team grouped by region
+        keys_by_region = get_team_keys_by_region(db, team.id)
 
         # Update keys for each region
         for region, keys in keys_by_region.items():
@@ -188,3 +241,220 @@ async def remove_product_from_team(db: Session, customer_id: str, product_id: st
         db.rollback()
         logger.error(f"Error removing product from team: {str(e)}")
         raise e
+
+async def monitor_teams(db: Session):
+    """
+    Daily monitoring task for teams that:
+    1. Posts age metrics for teams (since creation for teams without products, since last payment for teams with products)
+    2. Sends notifications for teams approaching expiration (25-30 days)
+    3. Posts metrics for expired teams (>30 days)
+    4. Monitors key spend and notifies if approaching limits
+    """
+    logger.info("Monitoring teams")
+    try:
+        # Initialize SES service
+        ses_service = SESService()
+
+        # Get all teams
+        teams = db.query(DBTeam).all()
+        current_time = datetime.now(UTC)
+
+        # Track current active team labels
+        current_team_labels = set()
+
+        logger.info(f"Found {len(teams)} teams to track")
+        for team in teams:
+            team_label = (str(team.id), team.name)
+            current_team_labels.add(team_label)
+
+            # Check if team has any products
+            has_products = db.query(DBTeamProduct).filter(
+                DBTeamProduct.team_id == team.id
+            ).first() is not None
+
+            # Calculate team age based on whether they have products
+            if has_products and team.last_payment:
+                team_freshness = (current_time - team.last_payment).days
+            else: # If a subscription is cancelled this will jump dramatically.
+                team_freshness = (current_time - team.created_at).days
+
+            if team_freshness < 0:
+                logger.warning(f"Team {team.name} (ID: {team.id}) has a negative age: {team_freshness} days")
+                team_freshness = 0
+
+            # Post freshness metric
+            team_freshness_days.labels(
+                team_id=str(team.id),
+                team_name=team.name
+            ).set(team_freshness)
+
+            # Check for notification conditions for teams without products
+            if not has_products:
+                days_remaining = TRIAL_OVER_DAYS - team_freshness
+                if days_remaining == FIRST_EMAIL_DAYS_LEFT or days_remaining == SECOND_EMAIL_DAYS_LEFT:
+                    logger.warning(f"Team {team.name} (ID: {team.id}) is approaching expiration in {days_remaining} days")
+                    # Send expiration notification email
+                    try:
+                        if team.admin_email:
+                            template_data = {
+                                "team_name": team.name,
+                                "days_remaining": days_remaining,
+                                "dashboard_url": generate_pricing_url(db, team)
+                            }
+                            ses_service.send_email(
+                                to_addresses=[team.admin_email],
+                                template_name="team-expiring",
+                                template_data=template_data
+                            )
+                            logger.info(f"Sent expiration notification email to team {team.name} (ID: {team.id})")
+                        else:
+                            logger.warning(f"No email found for team {team.name} (ID: {team.id})")
+                    except Exception as e:
+                        logger.error(f"Failed to send expiration notification email to team {team.name}: {str(e)}")
+                elif days_remaining <= 0:
+                    # Post expired metric
+                    team_expired_metric.labels(
+                        team_id=str(team.id),
+                        team_name=team.name
+                    ).inc()
+
+            # Get all keys for the team grouped by region
+            keys_by_region = get_team_keys_by_region(db, team.id)
+
+            # Track total spend across all keys for this team
+            team_total = 0
+
+            # Monitor keys for each region
+            for region, keys in keys_by_region.items():
+                try:
+                    # Initialize LiteLLM service for this region
+                    litellm_service = LiteLLMService(
+                        api_url=region.litellm_api_url,
+                        api_key=region.litellm_api_key
+                    )
+
+                    # Check spend for each key in this region
+                    for key in keys:
+                        try:
+                            # Get current spend using get_key_info
+                            key_info = await litellm_service.get_key_info(key.litellm_token)
+                            info = key_info.get("info", {})
+                            current_spend = info.get("spend", 0)
+                            budget = info.get("max_budget", 0)
+                            key_alias = info.get("key_alias", f"key-{key.id}")  # Fallback to key-{id} if no alias
+
+                            # Add to team total
+                            team_total += current_spend
+
+                            # Calculate and post percentage used
+                            if budget > 0:
+                                percentage_used = (current_spend / budget) * 100
+                                key_spend_percentage.labels(
+                                    team_id=str(team.id),
+                                    team_name=team.name,
+                                    key_alias=key_alias
+                                ).set(percentage_used)
+
+                                # Log warning if approaching limit
+                                if percentage_used >= 80:
+                                    logger.warning(
+                                        f"Key {key_alias} for team {team.name} is approaching spend limit: "
+                                        f"${current_spend:.2f} of ${budget:.2f} ({percentage_used:.1f}%)"
+                                    )
+                            else:
+                                # Set to 0 if no budget is set
+                                key_spend_percentage.labels(
+                                    team_id=str(team.id),
+                                    team_name=team.name,
+                                    key_alias=key_alias
+                                ).set(0)
+
+                        except Exception as e:
+                            logger.error(f"Error monitoring key {key.id} spend: {str(e)}")
+                            continue
+
+                except Exception as e:
+                    logger.error(f"Error initializing LiteLLM service for region {region.name}: {str(e)}")
+                    continue
+
+            # Set the total spend metric for the team
+            team_total_spend.labels(
+                team_id=str(team.id),
+                team_name=team.name
+            ).set(team_total)
+
+        # Zero out metrics for teams that are no longer active
+        for old_label in active_team_labels - current_team_labels:
+            team_freshness_days.labels(
+                team_id=old_label[0],
+                team_name=old_label[1]
+            ).set(0)
+
+        # Update active team labels for next run
+        active_team_labels.clear()
+        active_team_labels.update(current_team_labels)
+
+    except Exception as e:
+        logger.error(f"Error in team monitoring task: {str(e)}")
+        raise e
+
+def generate_team_admin_token(db: Session, team: DBTeam) -> str:
+    """
+    Generate a JWT token that authorizes the bearer as an administrator of the team.
+
+    Args:
+        db: Database session
+        team: The team object to generate the token for
+
+    Returns:
+        str: The generated JWT token
+
+    Raises:
+        ValueError: If no admin user is found for the team
+    """
+    token_validity_days = 1
+    # Find a team admin user
+    admin_user = db.query(DBUser).filter(
+        DBUser.team_id == team.id,
+        DBUser.role == "admin"
+    ).first()
+
+    if not admin_user:
+        raise ValueError(f"No admin user found for team {team.name} (ID: {team.id})")
+
+    # Create token payload with team admin claims
+    payload = {
+        "sub": admin_user.email,
+        "exp": datetime.now(UTC) + timedelta(days=token_validity_days)
+    }
+
+    # Generate the token
+    token = create_access_token(
+        data=payload,
+        expires_delta=timedelta(days=token_validity_days)
+    )
+
+    return token
+
+def generate_pricing_url(db: Session, team: DBTeam) -> str:
+    """
+    Generate a URL for the team admin pricing page with a JWT token.
+
+    Args:
+        db: Database session
+        team: The team object to generate the URL for
+
+    Returns:
+        str: The generated URL with the JWT token
+    """
+    # Generate the token
+    token = generate_team_admin_token(db, team)
+
+    # Get the frontend URL from settings
+    base_url = settings.frontend_route
+    path = '/pricing'
+    url = urljoin(base_url, path)
+
+    # Add the token as a query parameter
+    return f"{url}?token={token}"
+