Cache sales data in recon

app/api/teams.py

@@ -6,7 +6,7 @@
 import logging
 
 from app.db.database import get_db
-from app.db.models import DBTeam, DBTeamProduct, DBUser, DBPrivateAIKey, DBRegion, DBTeamRegion, DBProduct
+from app.db.models import DBTeam, DBTeamProduct, DBUser, DBPrivateAIKey, DBRegion, DBTeamRegion, DBProduct, DBTeamMetrics
 from app.core.security import get_role_min_system_admin, get_role_min_specific_team_admin, get_current_user_from_auth, check_sales_or_higher
 from app.schemas.models import (
     Team, TeamCreate, TeamUpdate,
@@ -248,14 +248,6 @@ async def list_teams_for_sales(
         all_regions = db.query(DBRegion).filter(DBRegion.is_active == True).all()
         regions_map = {r.id: r for r in all_regions}
 
-        # Pre-create LiteLLM services for each region to avoid re-instantiation
-        litellm_services = {}
-        for region in all_regions:
-            litellm_services[region.id] = LiteLLMService(
-                api_url=region.litellm_api_url,
-                api_key=region.litellm_api_key
-            )
-
         # Get all teams with their basic information
         teams = db.query(DBTeam).all()
 
@@ -277,41 +269,70 @@ async def list_teams_for_sales(
                 for team_product in team_products
             ]
 
-            # Get team AI keys (both team-owned and user-owned) and calculate total spend
-            team_users = db.query(DBUser).filter(DBUser.team_id == team.id).all()
-            team_user_ids = [user.id for user in team_users]
-
-            team_keys = db.query(DBPrivateAIKey).filter(
-                (DBPrivateAIKey.team_id == team.id) |  # Team-owned keys
-                (DBPrivateAIKey.owner_id.in_(team_user_ids))  # User-owned keys by team members
-            ).all()
+            # Try to get cached metrics first
+            team_metrics = db.query(DBTeamMetrics).filter(DBTeamMetrics.team_id == team.id).first()
+
+            if team_metrics:
+                # Use cached data
+                total_spend = team_metrics.total_spend
+                regions = team_metrics.regions or []
+            else:
+                # Fallback to real-time calculation if no cached data
+                logger.warning(f"No cached metrics found for team {team.id}, falling back to real-time calculation")
+                current_time = datetime.now(UTC)
+
+                # Create LiteLLM services only when needed for fallback
+                litellm_services = {}
+                for region in all_regions:
+                    litellm_services[region.id] = LiteLLMService(
+                        api_url=region.litellm_api_url,
+                        api_key=region.litellm_api_key
+                    )
 
-            # Calculate total spend from all AI keys and build regions list as we go
-            total_spend = 0.0
-            regions_set = set()
-
-            for key in team_keys:
-                if key.litellm_token and key.region_id in regions_map:
-                    try:
-                        # Use pre-fetched region info and pre-created LiteLLM service
-                        region = regions_map[key.region_id]
-                        litellm_service = litellm_services[region.id]
-
-                        # Add region name to our set
-                        regions_set.add(region.name)
-
-                        # Get spend data from LiteLLM
-                        key_data = await litellm_service.get_key_info(key.litellm_token)
-                        key_spend = key_data.get("info", {}).get("spend", 0.0)
-                        total_spend += float(key_spend)
-                    except Exception as e:
-                        # Track unreachable endpoint for logging at the end (only once per region)
-                        region = regions_map[key.region_id]
-                        endpoint_info = f"Region: {region.name}"
-                        unreachable_endpoints.add(endpoint_info)
-
-            # Convert set to list for the response
-            regions = list(regions_set)
+                # Get team AI keys (both team-owned and user-owned) and calculate total spend
+                team_users = db.query(DBUser).filter(DBUser.team_id == team.id).all()
+                team_user_ids = [user.id for user in team_users]
+
+                team_keys = db.query(DBPrivateAIKey).filter(
+                    (DBPrivateAIKey.team_id == team.id) |  # Team-owned keys
+                    (DBPrivateAIKey.owner_id.in_(team_user_ids))  # User-owned keys by team members
+                ).all()
+
+                # Calculate total spend from all AI keys and build regions list as we go
+                total_spend = 0.0
+                regions_set = set()
+
+                for key in team_keys:
+                    if key.litellm_token and key.region_id in regions_map:
+                        try:
+                            # Use pre-fetched region info and pre-created LiteLLM service
+                            region = regions_map[key.region_id]
+                            litellm_service = litellm_services[region.id]
+
+                            # Add region name to our set
+                            regions_set.add(region.name)
+
+                            # Get spend data from LiteLLM
+                            key_data = await litellm_service.get_key_info(key.litellm_token)
+                            key_spend = key_data.get("info", {}).get("spend", 0.0)
+                            total_spend += float(key_spend)
+                        except Exception as e:
+                            # Track unreachable endpoint for logging at the end (only once per region)
+                            region = regions_map[key.region_id]
+                            endpoint_info = f"Region: {region.name}"
+                            unreachable_endpoints.add(endpoint_info)
+
+                # Convert set to list for the response
+                regions = list(regions_set)
+                # Create new metrics record
+                team_metrics = DBTeamMetrics(
+                    team_id=team.id,
+                    total_spend=total_spend,
+                    last_spend_calculation=current_time,
+                    regions=regions,
+                    last_updated=current_time
+                )
+                db.add(team_metrics)
 
             # Calculate trial status
             trial_status = _calculate_trial_status(team, products)
@@ -330,6 +351,8 @@ async def list_teams_for_sales(
             )
 
             sales_teams.append(sales_team)
+        # Any metrics calculated on-the-fly will be cached
+        db.commit()
 
         # Log all unreachable endpoints at the end
         if unreachable_endpoints:

app/core/worker.py

@@ -1,7 +1,7 @@
 import re
 from datetime import datetime, UTC, timedelta
 from sqlalchemy.orm import Session
-from app.db.models import DBTeam, DBProduct, DBTeamProduct, DBPrivateAIKey, DBUser, DBRegion
+from app.db.models import DBTeam, DBProduct, DBTeamProduct, DBPrivateAIKey, DBUser, DBRegion, DBTeamMetrics
 from app.services.litellm import LiteLLMService
 from app.services.ses import SESService
 import logging
@@ -581,6 +581,31 @@ async def monitor_teams(db: Session):
                 team_name=team.name
             ).set(team_total)
 
+            # Update or create team metrics record
+            regions_list = list(keys_by_region.keys())
+            region_names = [region.name for region in regions_list]
+
+            # Check if metrics record exists
+            team_metrics = db.query(DBTeamMetrics).filter(DBTeamMetrics.team_id == team.id).first()
+
+            if team_metrics:
+                logger.info(f"metrics last updated at {team_metrics.last_updated}, curent time is {current_time}")
+                # Update existing metrics
+                team_metrics.total_spend = team_total
+                team_metrics.last_spend_calculation = current_time
+                team_metrics.regions = region_names
+                team_metrics.last_updated = current_time
+            else:
+                # Create new metrics record
+                team_metrics = DBTeamMetrics(
+                    team_id=team.id,
+                    total_spend=team_total,
+                    last_spend_calculation=current_time,
+                    regions=region_names,
+                    last_updated=current_time
+                )
+                db.add(team_metrics)
+
             # Update last_monitored timestamp only if notifications were sent
             if should_send_notifications:
                 team.last_monitored = current_time

app/db/models.py

@@ -18,7 +18,7 @@ class DBTeamProduct(Base):
     """
     __tablename__ = "team_products"
 
-    team_id = Column(Integer, ForeignKey('teams.id', ondelete='CASCADE'), primary_key=True, nullable=False)
+    team_id = Column(Integer, ForeignKey('teams.id'), primary_key=True, nullable=False)
     product_id = Column(String, ForeignKey('products.id', ondelete='CASCADE'), primary_key=True, nullable=False)
     created_at = Column(DateTime(timezone=True), default=func.now(), nullable=False)
     updated_at = Column(DateTime(timezone=True), onupdate=func.now(), nullable=True)
@@ -38,7 +38,7 @@ class DBTeamRegion(Base):
     """
     __tablename__ = "team_regions"
 
-    team_id = Column(Integer, ForeignKey('teams.id', ondelete='CASCADE'), primary_key=True, nullable=False)
+    team_id = Column(Integer, ForeignKey('teams.id'), primary_key=True, nullable=False)
     region_id = Column(Integer, ForeignKey('regions.id', ondelete='CASCADE'), primary_key=True, nullable=False)
     created_at = Column(DateTime(timezone=True), default=func.now(), nullable=False)
     updated_at = Column(DateTime(timezone=True), onupdate=func.now(), nullable=True)
@@ -114,6 +114,31 @@ class DBTeam(Base):
     private_ai_keys = relationship("DBPrivateAIKey", back_populates="team")
     active_products = relationship("DBTeamProduct", back_populates="team")
     dedicated_regions = relationship("DBTeamRegion", back_populates="team")
+    metrics = relationship("DBTeamMetrics", back_populates="team", uselist=False, cascade="all, delete")
+
+class DBTeamMetrics(Base):
+    """
+    Cached team metrics table populated by the monitor_teams worker.
+    This table stores pre-calculated metrics to avoid expensive real-time
+    LiteLLM API calls in the sales dashboard.
+    """
+    __tablename__ = "team_metrics"
+
+    id = Column(Integer, primary_key=True, index=True)
+    team_id = Column(Integer, ForeignKey('teams.id', ondelete='CASCADE'), unique=True, nullable=False, index=True)
+
+    # Spend metrics (the expensive calculation we want to cache)
+    total_spend = Column(Float, default=0.0, nullable=False)
+    last_spend_calculation = Column(DateTime(timezone=True), nullable=False)
+
+    # Region information (derived from team keys)
+    regions = Column(JSON, nullable=True)  # List of region names
+
+    # Monitoring metadata
+    last_updated = Column(DateTime(timezone=True), default=func.now(), nullable=False)
+
+    # Relationships
+    team = relationship("DBTeam", back_populates="metrics")
 
 class DBPrivateAIKey(Base):
     __tablename__ = "ai_tokens"

app/migrations/versions/20250127_120000_add_team_metrics_table.py

@@ -0,0 +1,45 @@
+"""add team metrics table
+
+Revision ID: add_team_metrics_table
+Revises: 5bda44ccd008
+Create Date: 2025-01-27 12:00:00.000000+00:00
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = 'add_team_metrics_table'
+down_revision: Union[str, None] = '5bda44ccd008'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table('team_metrics',
+    sa.Column('id', sa.Integer(), nullable=False),
+    sa.Column('team_id', sa.Integer(), nullable=False),
+    sa.Column('total_spend', sa.Float(), nullable=False),
+    sa.Column('last_spend_calculation', sa.DateTime(timezone=True), nullable=False),
+    sa.Column('regions', sa.JSON(), nullable=True),
+    sa.Column('last_updated', sa.DateTime(timezone=True), nullable=False),
+    sa.ForeignKeyConstraint(['team_id'], ['teams.id'], ondelete='CASCADE'),
+    sa.PrimaryKeyConstraint('id'),
+    sa.UniqueConstraint('team_id')
+    )
+    op.create_index(op.f('ix_team_metrics_id'), 'team_metrics', ['id'], unique=False)
+    op.create_index(op.f('ix_team_metrics_team_id'), 'team_metrics', ['team_id'], unique=False)
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_index(op.f('ix_team_metrics_team_id'), table_name='team_metrics')
+    op.drop_index(op.f('ix_team_metrics_id'), table_name='team_metrics')
+    op.drop_table('team_metrics')
+    # ### end Alembic commands ###
+