From 757abc1b09a9fd765674cdbf90826135bb456ee6 Mon Sep 17 00:00:00 2001 From: Pippa H Date: Wed, 13 Aug 2025 13:24:49 +0200 Subject: [PATCH 1/4] Enhance monitor_team_keys function to support renewal period updates - Added optional parameters for renewal_period_days and max_budget_amount to monitor_team_keys. - Implemented logic to check and update key budgets based on renewal periods and recent budget resets. - Updated tests to verify the new functionality, including scenarios for renewal period checks and budget updates based on heuristics. --- app/core/worker.py | 129 ++++++++++- tests/test_worker.py | 537 ++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 655 insertions(+), 11 deletions(-) diff --git a/app/core/worker.py b/app/core/worker.py index f8338d8..ac4275b 100644 --- a/app/core/worker.py +++ b/app/core/worker.py @@ -1,3 +1,4 @@ +import re from datetime import datetime, UTC, timedelta from sqlalchemy.orm import Session from app.db.models import DBTeam, DBProduct, DBTeamProduct, DBPrivateAIKey, DBUser, DBRegion @@ -17,7 +18,7 @@ INVOICE_SUCCESS_EVENTS ) from prometheus_client import Gauge, Counter, Summary -from typing import Dict, List +from typing import Dict, List, Optional from app.core.security import create_access_token from app.core.config import settings from urllib.parse import urljoin @@ -95,6 +96,8 @@ def get_team_keys_by_region(db: Session, team_id: int) -> Dict[DBRegion, List[DB logger.info(f"Found {len(team_keys)} keys in {len(keys_by_region)} regions for team {team_id}") return keys_by_region + + async def handle_stripe_event_background(event, db: Session): """ Background task to handle Stripe webhook events. @@ -247,19 +250,28 @@ async def remove_product_from_team(db: Session, customer_id: str, product_id: st logger.error(f"Error removing product from team: {str(e)}") raise e -async def monitor_team_keys(team: DBTeam, keys_by_region: Dict[DBRegion, List[DBPrivateAIKey]], expire_keys: bool) -> float: +async def monitor_team_keys( + team: DBTeam, + keys_by_region: Dict[DBRegion, List[DBPrivateAIKey]], + expire_keys: bool, + renewal_period_days: Optional[int] = None, + max_budget_amount: Optional[float] = None +) -> float: """ - Monitor spend for all keys in a team across different regions. + Monitor spend for all keys in a team across different regions and optionally update keys after renewal period. Args: - db: Database session team: The team to monitor keys for keys_by_region: Dictionary mapping regions to lists of keys + expire_keys: Whether to expire keys (set duration to 0) + renewal_period_days: Optional renewal period in days. If provided, will check for and update keys renewed within the last hour. + max_budget_amount: Optional maximum budget amount. If provided, will update the budget amount for the keys. Returns: float: Total spend across all keys for the team """ team_total = 0 + current_time = datetime.now(UTC) # Monitor keys for each region for region, keys in keys_by_region.items(): @@ -281,6 +293,86 @@ async def monitor_team_keys(team: DBTeam, keys_by_region: Dict[DBRegion, List[DB budget = info.get("max_budget", 0) or 0.0 key_alias = info.get("key_alias", f"key-{key.id}") # Fallback to key-{id} if no alias + # Check for renewal period update if renewal_period_days is provided + if renewal_period_days is not None: + # Check current values and only update if they don't match the parameters + current_budget_duration = info.get("budget_duration") + current_max_budget = info.get("max_budget") + + needs_update = False + data = {"litellm_token": key.litellm_token} + + # Check if budget_duration needs updating + expected_budget_duration = f"{renewal_period_days}d" + if current_budget_duration != expected_budget_duration: + data["budget_duration"] = expected_budget_duration + needs_update = True + logger.info(f"Key {key.id} budget_duration will be updated from '{current_budget_duration}' to '{expected_budget_duration}'") + + # Check if budget_amount needs updating + if max_budget_amount is not None and current_max_budget != max_budget_amount: + data["budget_amount"] = max_budget_amount + needs_update = True + logger.info(f"Key {key.id} budget_amount will be updated from {current_max_budget} to {max_budget_amount}") + + # Only check reset timestamp if we need to update + if needs_update: + budget_reset_at_str = info.get("budget_reset_at") + if budget_reset_at_str: + try: + # Parse the budget_reset_at timestamp + budget_reset_at = datetime.fromisoformat(budget_reset_at_str.replace('Z', '+00:00')) + if budget_reset_at.tzinfo is None: + budget_reset_at = budget_reset_at.replace(tzinfo=UTC) + logger.info(f"Key {key.id} budget_reset_at_str: {budget_reset_at_str}, budget_reset_at: {budget_reset_at}") + + # Check if budget was reset recently using heuristics + # budget_reset_at represents when the next reset will occur + current_spend = info.get("spend", 0) or 0.0 + current_budget_duration = info.get("budget_duration") + + should_update = False + update_reason = "" + + # Heuristic 1: Check if (now + current_budget_duration) is within an hour of budget_reset_at + if current_budget_duration is not None: + try: + # Parse current budget duration (e.g., "30d" -> 30 days) + duration_match = re.match(r'(\d+)d', current_budget_duration) + if duration_match: + duration_days = int(duration_match.group(1)) + expected_reset_time = current_time + timedelta(days=duration_days) + hours_diff = abs((expected_reset_time - budget_reset_at).total_seconds() / 3600) + + if hours_diff <= 1.0: + should_update = True + update_reason = f"reset time alignment (within {hours_diff:.2f} hours)" + except (ValueError, AttributeError): + logger.warning(f"Key {key.id} has invalid budget_duration format: {current_budget_duration}") + else: + logger.debug(f"Key {key.id} has no budget_duration set, skipping reset time alignment heuristic") + should_update = True + update_reason = "no budget_duration set, forcing update" + + # Heuristic 2: Update if amount spent is $0.00 (indicating fresh reset) + if current_spend == 0.0: + should_update = True + update_reason = "zero spend (fresh reset)" + + if should_update: + logger.info(f"Key {key.id} budget update triggered: {update_reason}, updating budget settings") + await litellm_service.update_budget(**data) + logger.info(f"Updated key {key.id} budget settings") + else: + logger.debug(f"Key {key.id} budget update not triggered, skipping update") + except ValueError: + logger.warning(f"Key {key.id} has invalid budget_reset_at timestamp: {budget_reset_at_str}") + else: + logger.warning(f"Key {key.id} has no budget_reset_at timestamp, forcing update") + await litellm_service.update_budget(**data) + else: + logger.info(f"Key {key.id} budget settings already match the expected values, no update needed") + # Set the key duration to 0 days to end its usability. if expire_keys: await litellm_service.update_key_duration(key.litellm_token, "0d") @@ -332,15 +424,18 @@ async def monitor_teams(db: Session): """ 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() + try: + # Initialize SES service + ses_service = SESService() + except Exception as e: + logger.error(f"Error initializing SES service: {str(e)}") + pass logger.info(f"Found {len(teams)} teams to track") for team in teams: @@ -436,8 +531,24 @@ async def monitor_teams(db: Session): if not has_products and days_remaining <= 0 and should_send_notifications: expire_keys = True - # Monitor keys and get total spend - team_total = await monitor_team_keys(team, keys_by_region, expire_keys) + # Determine if we should check for renewal period updates + renewal_period_days = None + max_budget_amount = None + if has_products and team.last_payment: + # Get the product with the longest renewal period + active_products = db.query(DBTeamProduct).filter( + DBTeamProduct.team_id == team.id + ).all() + product_ids = [tp.product_id for tp in active_products] + products = db.query(DBProduct).filter(DBProduct.id.in_(product_ids)).all() + + if products: + max_renewal_product = max(products, key=lambda product: product.renewal_period_days) + renewal_period_days = max_renewal_product.renewal_period_days + max_budget_amount = max(products, key=lambda product: product.max_budget_per_key).max_budget_per_key + + # Monitor keys and get total spend (includes renewal period updates if applicable) + team_total = await monitor_team_keys(team, keys_by_region, expire_keys, renewal_period_days, max_budget_amount) # Set the total spend metric for the team (always emit metrics) team_total_spend.labels( diff --git a/tests/test_worker.py b/tests/test_worker.py index b97ab95..abaa5bc 100644 --- a/tests/test_worker.py +++ b/tests/test_worker.py @@ -10,7 +10,9 @@ team_expired_metric, key_spend_percentage, team_total_spend, - active_team_labels + active_team_labels, + get_team_keys_by_region, + monitor_team_keys ) from unittest.mock import AsyncMock, patch, Mock @@ -545,6 +547,10 @@ async def test_monitor_teams_basic_metrics(mock_litellm, mock_ses, db, test_team db.add(team_product) db.commit() + # Setup mock LiteLLM service + mock_instance = mock_litellm.return_value + mock_instance.get_key_info = AsyncMock(return_value={"info": {"spend": 0, "max_budget": 100, "key_alias": "test"}}) + # Run monitoring await monitor_teams(db) @@ -941,4 +947,531 @@ async def test_monitor_teams_metrics_always_emitted(mock_litellm, mock_ses, db, # Verify last_monitored was not updated (no notifications sent) db.refresh(test_team) # Use approximate comparison due to timestamp precision differences - assert abs((test_team.last_monitored - expected_last_monitored).total_seconds()) < 1 \ No newline at end of file + assert abs((test_team.last_monitored - expected_last_monitored).total_seconds()) < 1 + +@pytest.mark.asyncio +@patch('app.core.worker.SESService') +@patch('app.core.worker.LiteLLMService') +@patch('app.core.config.settings.ENABLE_LIMITS', True) +async def test_monitor_teams_includes_renewal_period_check(mock_litellm, mock_ses, db, test_team, test_product, test_region): + """ + Test that the monitoring workflow includes renewal period checks when conditions are met. + + Given: A team with an active product that has passed its renewal period + When: The monitoring workflow runs + Then: The monitor_team_keys function should be called with renewal_period_days + """ + # Setup test data + test_team.last_payment = datetime.now(UTC) - timedelta(days=35) # 35 days ago (past 30-day renewal period) + db.add(test_team) + + # Add product to team + team_product = DBTeamProduct( + team_id=test_team.id, + product_id=test_product.id + ) + db.add(team_product) + + # Create a key for the team + team_key = DBPrivateAIKey( + name="Team Key", + litellm_token="team_token_123", + region=test_region, + team_id=test_team.id + ) + db.add(team_key) + db.commit() + + # Setup mocks + mock_instance = mock_litellm.return_value + mock_instance.get_key_info = AsyncMock(return_value={"info": {"spend": 0, "max_budget": 100, "key_alias": "test"}}) + + # Run monitoring + await monitor_teams(db) + + # Verify that get_key_info was called (indicating the combined function ran) + # The function should have been called to get key info for monitoring AND renewal period checks + assert mock_instance.get_key_info.called + +@pytest.mark.asyncio +@patch('app.core.worker.SESService') +@patch('app.core.worker.LiteLLMService') +@patch('app.core.config.settings.ENABLE_LIMITS', True) +async def test_monitor_teams_does_not_include_renewal_period_check_when_not_passed(mock_litellm, mock_ses, db, test_team, test_product, test_region): + """ + Test that the monitoring workflow does not include renewal period checks when conditions are not met. + + Given: A team with an active product but renewal period hasn't passed + When: The monitoring workflow runs + Then: The monitor_team_keys function should be called without renewal_period_days + """ + # Setup test data + test_team.last_payment = datetime.now(UTC) - timedelta(days=15) # 15 days ago (before 30-day renewal period) + db.add(test_team) + + # Add product to team + team_product = DBTeamProduct( + team_id=test_team.id, + product_id=test_product.id + ) + db.add(team_product) + + # Create a key for the team + team_key = DBPrivateAIKey( + name="Team Key", + litellm_token="team_token_123", + region=test_region, + team_id=test_team.id + ) + db.add(team_key) + db.commit() + + # Setup mocks + mock_instance = mock_litellm.return_value + mock_instance.get_key_info = AsyncMock(return_value={"info": {"spend": 0, "max_budget": 100, "key_alias": "test"}}) + + # Run monitoring + await monitor_teams(db) + + # Verify that get_key_info was called (for monitoring) but no renewal period updates occurred + # Since renewal period hasn't passed, the function should still be called but without renewal checks + assert mock_instance.get_key_info.called + +@pytest.mark.asyncio +@patch('app.core.worker.LiteLLMService') +async def test_monitor_team_keys_with_renewal_period_updates(mock_litellm, db, test_team, test_product, test_region, test_team_user, test_team_key_creator): + """ + Test that monitor_team_keys updates keys after renewal period when LiteLLM has reset their budget within the last hour. + + Given: A team with keys that have had their budget reset within the last hour + When: monitor_team_keys is called with renewal_period_days + Then: The budget_duration should be updated to match the product renewal period + """ + # Setup test data + test_team.last_payment = datetime.now(UTC) - timedelta(days=35) # 35 days ago (past 30-day renewal period) + db.add(test_team) + + # Add product to team + team_product = DBTeamProduct( + team_id=test_team.id, + product_id=test_product.id + ) + db.add(team_product) + + # Create keys for the team + team_key = DBPrivateAIKey( + name="Team Key", + litellm_token="team_token_123", + region=test_region, + team_id=test_team.id + ) + db.add(team_key) + + user_key = DBPrivateAIKey( + name="User Key", + litellm_token="user_token_456", + region=test_region, + owner_id=test_team_user.id + ) + db.add(user_key) + + db.commit() + + # Setup mock LiteLLM service + mock_instance = mock_litellm.return_value + mock_instance.get_key_info = AsyncMock() + mock_instance.update_budget = AsyncMock() + + # Mock key info responses - one key with zero spend (triggers heuristic), one not + mock_instance.get_key_info.side_effect = [ + # Team key - zero spend (triggers zero spend heuristic) + { + "info": { + "budget_reset_at": (datetime.now(UTC) + timedelta(days=30)).isoformat(), + "key_alias": "team_key", + "spend": 0.0, # Zero spend triggers update + "max_budget": 100.0, + "budget_duration": "15d" + } + }, + # User key - non-zero spend (doesn't trigger heuristic) + { + "info": { + "budget_reset_at": (datetime.now(UTC) + timedelta(days=30)).isoformat(), + "key_alias": "user_key", + "spend": 5.0, # Non-zero spend + "max_budget": 50.0, + "budget_duration": "15d" + } + } + ] + + # Get keys by region + keys_by_region = get_team_keys_by_region(db, test_team.id) + + # Call the combined function with renewal period days and budget amount + team_total = await monitor_team_keys(test_team, keys_by_region, False, test_product.renewal_period_days, test_product.max_budget_per_key) + + # Verify LiteLLM service was initialized correctly + mock_litellm.assert_called_once_with( + api_url=test_region.litellm_api_url, + api_key=test_region.litellm_api_key + ) + + # Verify get_key_info was called for both keys + assert mock_instance.get_key_info.call_count == 2 + + # Verify update_budget was called only for the key with budget reset within the last hour + assert mock_instance.update_budget.call_count == 1 + update_call = mock_instance.update_budget.call_args + assert update_call[1]['litellm_token'] == "team_token_123" + assert update_call[1]['budget_duration'] == f"{test_product.renewal_period_days}d" + assert update_call[1]['budget_amount'] == test_product.max_budget_per_key + + # Verify team total spend is calculated correctly + assert team_total == 5.0 # 0.0 + 5.0 + +@pytest.mark.asyncio +@patch('app.core.worker.LiteLLMService') +async def test_monitor_team_keys_with_renewal_period_updates_no_products(mock_litellm, db, test_team, test_region, test_team_user, test_team_key_creator): + """ + Test that monitor_team_keys updates budget_duration even when no products are found. + + Given: A team with keys that have had their budget reset within the last hour, but no active products + When: monitor_team_keys is called with renewal_period_days + Then: The budget_duration should be updated but budget_amount should not be set + """ + # Setup test data - team with no products + test_team.last_payment = datetime.now(UTC) - timedelta(days=35) # 35 days ago (past 30-day renewal period) + db.add(test_team) + + # Create keys for the team + team_key = DBPrivateAIKey( + name="Team Key", + litellm_token="team_token_123", + region=test_region, + team_id=test_team.id + ) + db.add(team_key) + + user_key = DBPrivateAIKey( + name="User Key", + litellm_token="user_token_456", + region=test_region, + owner_id=test_team_user.id + ) + db.add(user_key) + + db.commit() + + # Setup mock LiteLLM service + mock_instance = mock_litellm.return_value + mock_instance.get_key_info = AsyncMock() + mock_instance.update_budget = AsyncMock() + + # Mock key info responses - one key with zero spend (triggers heuristic), one not + mock_instance.get_key_info.side_effect = [ + # Team key - zero spend (triggers zero spend heuristic) + { + "info": { + "budget_reset_at": (datetime.now(UTC) + timedelta(days=30)).isoformat(), + "key_alias": "team_key", + "spend": 0.0, # Zero spend triggers update + "max_budget": 100.0, + "budget_duration": "15d" + } + }, + # User key - non-zero spend (doesn't trigger heuristic) + { + "info": { + "budget_reset_at": (datetime.now(UTC) + timedelta(days=30)).isoformat(), + "key_alias": "user_key", + "spend": 5.0, # Non-zero spend + "max_budget": 50.0, + "budget_duration": "15d" + } + } + ] + + # Get keys by region + keys_by_region = get_team_keys_by_region(db, test_team.id) + + # Call the combined function with renewal period days (no budget amount) + team_total = await monitor_team_keys(test_team, keys_by_region, False, 30, None) # Use default 30 days, no budget amount + + # Verify LiteLLM service was initialized correctly + mock_litellm.assert_called_once_with( + api_url=test_region.litellm_api_url, + api_key=test_region.litellm_api_key + ) + + # Verify get_key_info was called for both keys + assert mock_instance.get_key_info.call_count == 2 + + # Verify update_budget was called only for the key with budget reset within the last hour + assert mock_instance.update_budget.call_count == 1 + update_call = mock_instance.update_budget.call_args + assert update_call[1]['litellm_token'] == "team_token_123" + assert update_call[1]['budget_duration'] == "30d" + # Should not have budget_amount since no products were found + assert 'budget_amount' not in update_call[1] + + # Verify team total spend is calculated correctly + assert team_total == 5.0 # 0.0 + 5.0 + +@pytest.mark.asyncio +@patch('app.core.worker.LiteLLMService') +async def test_monitor_team_keys_reset_time_alignment_heuristic(mock_litellm, db, test_team, test_product, test_region, test_team_user, test_team_key_creator): + """ + Test that monitor_team_keys updates when reset time alignment heuristic is met. + + Given: A team with keys where (now + current_budget_duration) is within 1 hour of budget_reset_at + When: monitor_team_keys is called with renewal_period_days + Then: The budget should be updated due to reset time alignment + """ + # Setup test data + test_team.last_payment = datetime.now(UTC) - timedelta(days=35) # 35 days ago (past 30-day renewal period) + db.add(test_team) + + # Add product to team + team_product = DBTeamProduct( + team_id=test_team.id, + product_id=test_product.id + ) + db.add(team_product) + + # Create a key for the team + team_key = DBPrivateAIKey( + name="Team Key", + litellm_token="team_token_123", + region=test_region, + team_id=test_team.id + ) + db.add(team_key) + db.commit() + + # Setup mock LiteLLM service + mock_instance = mock_litellm.return_value + mock_instance.get_key_info = AsyncMock() + mock_instance.update_budget = AsyncMock() + + current_time = datetime.now(UTC) + + # Mock key info response - current budget_duration is 15d, budget_reset_at is 30 minutes after (now + 15d) + # This should trigger the reset time alignment heuristic + mock_instance.get_key_info.return_value = { + "info": { + "budget_reset_at": (current_time + timedelta(days=15, minutes=30)).isoformat(), + "key_alias": "team_key", + "spend": 10.0, + "max_budget": 100.0, + "budget_duration": "15d" # Different from expected + } + } + + # Get keys by region + keys_by_region = get_team_keys_by_region(db, test_team.id) + + # Call the function with renewal period days + team_total = await monitor_team_keys(test_team, keys_by_region, False, test_product.renewal_period_days, test_product.max_budget_per_key) + + # Verify update_budget was called because of reset time alignment + assert mock_instance.update_budget.call_count == 1 + update_call = mock_instance.update_budget.call_args + assert update_call[1]['litellm_token'] == "team_token_123" + assert update_call[1]['budget_duration'] == f"{test_product.renewal_period_days}d" + assert update_call[1]['budget_amount'] == test_product.max_budget_per_key + + # Verify team total spend is calculated correctly + assert team_total == 10.0 + +@pytest.mark.asyncio +@patch('app.core.worker.LiteLLMService') +async def test_monitor_team_keys_zero_spend_heuristic(mock_litellm, db, test_team, test_product, test_region, test_team_user, test_team_key_creator): + """ + Test that monitor_team_keys updates when zero spend heuristic is met. + + Given: A team with keys where spend is $0.00 (indicating fresh reset) + When: monitor_team_keys is called with renewal_period_days + Then: The budget should be updated due to zero spend + """ + # Setup test data + test_team.last_payment = datetime.now(UTC) - timedelta(days=35) # 35 days ago (past 30-day renewal period) + db.add(test_team) + + # Add product to team + team_product = DBTeamProduct( + team_id=test_team.id, + product_id=test_product.id + ) + db.add(team_product) + + # Create a key for the team + team_key = DBPrivateAIKey( + name="Team Key", + litellm_token="team_token_123", + region=test_region, + team_id=test_team.id + ) + db.add(team_key) + db.commit() + + # Setup mock LiteLLM service + mock_instance = mock_litellm.return_value + mock_instance.get_key_info = AsyncMock() + mock_instance.update_budget = AsyncMock() + + current_time = datetime.now(UTC) + + # Mock key info response - spend is 0.0 (should trigger zero spend heuristic) + mock_instance.get_key_info.return_value = { + "info": { + "budget_reset_at": (current_time + timedelta(days=30)).isoformat(), + "key_alias": "team_key", + "spend": 0.0, # Zero spend triggers update + "max_budget": 100.0, + "budget_duration": "15d" # Different from expected + } + } + + # Get keys by region + keys_by_region = get_team_keys_by_region(db, test_team.id) + + # Call the function with renewal period days + team_total = await monitor_team_keys(test_team, keys_by_region, False, test_product.renewal_period_days, test_product.max_budget_per_key) + + # Verify update_budget was called because of zero spend + assert mock_instance.update_budget.call_count == 1 + update_call = mock_instance.update_budget.call_args + assert update_call[1]['litellm_token'] == "team_token_123" + assert update_call[1]['budget_duration'] == f"{test_product.renewal_period_days}d" + assert update_call[1]['budget_amount'] == test_product.max_budget_per_key + + # Verify team total spend is calculated correctly + assert team_total == 0.0 + +@pytest.mark.asyncio +@patch('app.core.worker.LiteLLMService') +async def test_monitor_team_keys_no_heuristic_met(mock_litellm, db, test_team, test_product, test_region, test_team_user, test_team_key_creator): + """ + Test that monitor_team_keys does not update when no heuristics are met. + + Given: A team with keys where neither reset time alignment nor zero spend heuristics are met + When: monitor_team_keys is called with renewal_period_days + Then: The budget should not be updated + """ + # Setup test data + test_team.last_payment = datetime.now(UTC) - timedelta(days=35) # 35 days ago (past 30-day renewal period) + db.add(test_team) + + # Add product to team + team_product = DBTeamProduct( + team_id=test_team.id, + product_id=test_product.id + ) + db.add(team_product) + + # Create a key for the team + team_key = DBPrivateAIKey( + name="Team Key", + litellm_token="team_token_123", + region=test_region, + team_id=test_team.id + ) + db.add(team_key) + db.commit() + + # Setup mock LiteLLM service + mock_instance = mock_litellm.return_value + mock_instance.get_key_info = AsyncMock() + mock_instance.update_budget = AsyncMock() + + current_time = datetime.now(UTC) + + # Mock key info response - no heuristics met + mock_instance.get_key_info.return_value = { + "info": { + "budget_reset_at": (current_time + timedelta(days=30)).isoformat(), + "key_alias": "team_key", + "spend": 10.0, # Non-zero spend + "max_budget": 100.0, + "budget_duration": "15d" # Different from expected, but reset time not aligned + } + } + + # Get keys by region + keys_by_region = get_team_keys_by_region(db, test_team.id) + + # Call the function with renewal period days + team_total = await monitor_team_keys(test_team, keys_by_region, False, test_product.renewal_period_days, test_product.max_budget_per_key) + + # Verify update_budget was NOT called because no heuristics are met + assert mock_instance.update_budget.call_count == 0 + + # Verify team total spend is calculated correctly + assert team_total == 10.0 + +@pytest.mark.asyncio +@patch('app.core.worker.LiteLLMService') +async def test_monitor_team_keys_none_budget_duration_handled(mock_litellm, db, test_team, test_product, test_region, test_team_user, test_team_key_creator): + """ + Test that monitor_team_keys handles None budget_duration gracefully. + + Given: A team with keys where budget_duration is None + When: monitor_team_keys is called with renewal_period_days + Then: The function should not error and should skip the reset time alignment heuristic + """ + # Setup test data + test_team.last_payment = datetime.now(UTC) - timedelta(days=35) # 35 days ago (past 30-day renewal period) + db.add(test_team) + + # Add product to team + team_product = DBTeamProduct( + team_id=test_team.id, + product_id=test_product.id + ) + db.add(team_product) + + # Create a key for the team + team_key = DBPrivateAIKey( + name="Team Key", + litellm_token="team_token_123", + region=test_region, + team_id=test_team.id + ) + db.add(team_key) + db.commit() + + # Setup mock LiteLLM service + mock_instance = mock_litellm.return_value + mock_instance.get_key_info = AsyncMock() + mock_instance.update_budget = AsyncMock() + + current_time = datetime.now(UTC) + + # Mock key info response - budget_duration is None, but spend is non-zero + mock_instance.get_key_info.return_value = { + "info": { + "budget_reset_at": (current_time + timedelta(days=30)).isoformat(), + "key_alias": "team_key", + "spend": 10.0, # Non-zero spend + "max_budget": 100.0, + "budget_duration": None # None budget_duration + } + } + + # Get keys by region + keys_by_region = get_team_keys_by_region(db, test_team.id) + + # Call the function with renewal period days + team_total = await monitor_team_keys(test_team, keys_by_region, False, test_product.renewal_period_days, test_product.max_budget_per_key) + + # Verify update_budget was called because budget_duration is None (forces update) + assert mock_instance.update_budget.call_count == 1 + update_call = mock_instance.update_budget.call_args + assert update_call[1]['litellm_token'] == "team_token_123" + assert update_call[1]['budget_duration'] == f"{test_product.renewal_period_days}d" + assert update_call[1]['budget_amount'] == test_product.max_budget_per_key + + # Verify team total spend is calculated correctly + assert team_total == 10.0 From 76ab41b714191ea0e70fb4891b033c3094a31b6b Mon Sep 17 00:00:00 2001 From: Pippa H Date: Wed, 13 Aug 2025 13:35:36 +0200 Subject: [PATCH 2/4] Refactor private AI keys data fetching to use individual queries - Changed the implementation to create separate queries for each private AI key's spend data. - Ensured stable hook count by using the keys array length. - Combined all spend data into a single map for easier access and management. --- .../src/hooks/use-private-ai-keys-data.ts | 32 ++++++++++++------- 1 file changed, 20 insertions(+), 12 deletions(-) diff --git a/frontend/src/hooks/use-private-ai-keys-data.ts b/frontend/src/hooks/use-private-ai-keys-data.ts index b161277..2bed9fb 100644 --- a/frontend/src/hooks/use-private-ai-keys-data.ts +++ b/frontend/src/hooks/use-private-ai-keys-data.ts @@ -62,18 +62,26 @@ export function usePrivateAIKeysData( }, }); - // Query to get spend information for each key - const { data: spendMap = {} } = useQuery>({ - queryKey: ['private-ai-keys-spend', Array.from(loadedSpendKeys)], - queryFn: async () => { - const spendPromises = Array.from(loadedSpendKeys).map(async (keyId) => { - const response = await get(`private-ai-keys/${keyId}/spend`); - return [keyId, await response.json()] as [number, SpendInfo]; - }); - const spendResults = await Promise.all(spendPromises); - return Object.fromEntries(spendResults); - }, - enabled: loadedSpendKeys.size > 0, + // Create individual queries for each key + // Use the keys array length to ensure stable hook count + const spendQueries = keys.map((key, index) => + useQuery({ + queryKey: ['private-ai-key-spend', key.id], + queryFn: async () => { + const response = await get(`private-ai-keys/${key.id}/spend`); + return response.json(); + }, + enabled: loadedSpendKeys.has(key.id), + }) + ); + + // Combine all spend data into a single map + const spendMap: Record = {}; + keys.forEach((key, index) => { + const query = spendQueries[index]; + if (query.data && loadedSpendKeys.has(key.id)) { + spendMap[key.id] = query.data; + } }); // Fetch regions From 4712666ee2c4145a21ce9d7ab07d115c01f26b08 Mon Sep 17 00:00:00 2001 From: Pippa H Date: Wed, 13 Aug 2025 13:53:44 +0200 Subject: [PATCH 3/4] Refactor private AI keys management and enhance spend data handling - Introduced a new component for displaying spend information for private AI keys. - Updated the private AI keys table to utilize the new spend cell component. - Simplified the data fetching logic for spend information by using a single query for loaded keys. - Removed unnecessary state management related to spend data loading. - Improved overall code organization and readability. --- .../src/app/admin/private-ai-keys/page.tsx | 94 ++++----- frontend/src/app/private-ai-keys/page.tsx | 18 +- .../app/team-admin/private-ai-keys/page.tsx | 62 +++--- .../components/private-ai-key-spend-cell.tsx | 192 ++++++++++++++++++ .../src/components/private-ai-keys-table.tsx | 106 +--------- .../src/hooks/use-private-ai-keys-data.ts | 34 ++-- 6 files changed, 286 insertions(+), 220 deletions(-) create mode 100644 frontend/src/components/private-ai-key-spend-cell.tsx diff --git a/frontend/src/app/admin/private-ai-keys/page.tsx b/frontend/src/app/admin/private-ai-keys/page.tsx index 307b508..f7df07f 100644 --- a/frontend/src/app/admin/private-ai-keys/page.tsx +++ b/frontend/src/app/admin/private-ai-keys/page.tsx @@ -1,6 +1,6 @@ 'use client'; -import { useState } from 'react'; +import { useState, useCallback } from 'react'; import { useQuery, useMutation, useQueryClient } from '@tanstack/react-query'; import { Button } from '@/components/ui/button'; import { useToast } from '@/hooks/use-toast'; @@ -34,19 +34,16 @@ interface User { created_at: string; } - - export default function PrivateAIKeysPage() { const { toast } = useToast(); const queryClient = useQueryClient(); - const [searchTerm, setSearchTerm] = useState(''); - const [isUserSearchOpen, setIsUserSearchOpen] = useState(false); - const [selectedUser, setSelectedUser] = useState(null); - const [loadedSpendKeys, setLoadedSpendKeys] = useState>(new Set()); const [isCreateDialogOpen, setIsCreateDialogOpen] = useState(false); + const [selectedUser, setSelectedUser] = useState(null); + const [isUserSearchOpen, setIsUserSearchOpen] = useState(false); + const [searchTerm, setSearchTerm] = useState(''); + const debouncedSearchTerm = useDebounce(searchTerm, 300); - // Queries const { data: privateAIKeys = [], isLoading: isLoadingPrivateAIKeys } = useQuery({ queryKey: ['private-ai-keys', selectedUser?.id], queryFn: async () => { @@ -56,15 +53,24 @@ export default function PrivateAIKeysPage() { const response = await get(url); const data = await response.json(); return data; - }, - refetchInterval: 30000, // Refetch every 30 seconds to detect new keys - refetchIntervalInBackground: true, // Continue polling even when tab is not active + } }); - // Use shared hook for data fetching - const { teamDetails, teamMembers, spendMap, regions } = usePrivateAIKeysData(privateAIKeys, loadedSpendKeys); + // Use shared hook for data fetching (only for team details and regions) + const { teamDetails, teamMembers, regions } = usePrivateAIKeysData(privateAIKeys, new Set()); - // Query to get all users for displaying emails + // Search users + const { data: users = [], isLoading: isSearching } = useQuery({ + queryKey: ['users', debouncedSearchTerm], + queryFn: async () => { + if (!debouncedSearchTerm) return []; + const response = await get(`/users?search=${debouncedSearchTerm}`); + return response.json(); + }, + enabled: debouncedSearchTerm.length > 0, + }); + + // Get all users for the dropdown const { data: usersMap = {} } = useQuery>({ queryKey: ['users-map'], queryFn: async () => { @@ -77,50 +83,22 @@ export default function PrivateAIKeysPage() { }, }); - const { data: users = [], isLoading: isLoadingUsers, isFetching: isFetchingUsers } = useQuery({ - queryKey: ['users', debouncedSearchTerm], - queryFn: async () => { - if (!debouncedSearchTerm) return []; - await new Promise(resolve => setTimeout(resolve, 100)); // Small delay to ensure loading state shows - const response = await get(`/users/search?email=${encodeURIComponent(debouncedSearchTerm)}`); - const data = await response.json(); - return data; - }, - enabled: isUserSearchOpen && !!debouncedSearchTerm, - gcTime: 60000, - staleTime: 30000, - refetchOnMount: false, - refetchOnWindowFocus: false, - }); - - // Show loading state immediately when search term changes - const isSearching = searchTerm.length > 0 && ( - isLoadingUsers || - isFetchingUsers || - debouncedSearchTerm !== searchTerm - ); - const handleSearchChange = (value: string) => { setSearchTerm(value); - // Prefetch the query if we have a value - if (value) { - queryClient.prefetchQuery({ - queryKey: ['users', value], - queryFn: async () => { - const response = await get(`/users/search?email=${encodeURIComponent(value)}`); - const data = await response.json(); - return data; - }, - }); - } }; - // Mutations + // Create key mutation const createKeyMutation = useMutation({ - mutationFn: async (data: { name: string; region_id: number; owner_id?: number; team_id?: number; key_type: 'full' | 'llm' | 'vector' }) => { - const endpoint = data.key_type === 'full' ? 'private-ai-keys' : - data.key_type === 'llm' ? 'private-ai-keys/token' : - 'private-ai-keys/vector-db'; + mutationFn: async (data: { + name: string + region_id: number + key_type: 'full' | 'llm' | 'vector' + owner_id?: number + team_id?: number + }) => { + const endpoint = data.key_type === 'full' ? '/private-ai-keys' : + data.key_type === 'llm' ? '/private-ai-keys/token' : + '/private-ai-keys/vector-db'; const response = await post(endpoint, data); return response.json(); }, @@ -142,9 +120,11 @@ export default function PrivateAIKeysPage() { }, }); + // Delete key mutation const deletePrivateAIKeyMutation = useMutation({ mutationFn: async (keyId: number) => { - await del(`/private-ai-keys/${keyId}`); + const response = await del(`/private-ai-keys/${keyId}`); + return response.json(); }, onSuccess: () => { queryClient.invalidateQueries({ queryKey: ['private-ai-keys'] }); @@ -172,8 +152,8 @@ export default function PrivateAIKeysPage() { return response.json(); }, onSuccess: (data, variables) => { - // Update the spend information for this specific key - queryClient.setQueryData(['private-ai-key-spend', variables.keyId], data); + // Invalidate the specific key's spend query to refresh the data + queryClient.invalidateQueries({ queryKey: ['private-ai-key-spend', variables.keyId] }); toast({ title: 'Success', description: 'Budget period updated successfully', @@ -284,8 +264,6 @@ export default function PrivateAIKeysPage() { isDeleting={deletePrivateAIKeyMutation.isPending} allowModification={true} showOwner={true} - spendMap={spendMap} - onLoadSpend={(keyId) => setLoadedSpendKeys(prev => new Set([...prev, keyId]))} onUpdateBudget={(keyId, budgetDuration) => { updateBudgetPeriodMutation.mutate({ keyId, budgetDuration }); }} diff --git a/frontend/src/app/private-ai-keys/page.tsx b/frontend/src/app/private-ai-keys/page.tsx index 0996e35..b5e5d27 100644 --- a/frontend/src/app/private-ai-keys/page.tsx +++ b/frontend/src/app/private-ai-keys/page.tsx @@ -32,7 +32,6 @@ export default function DashboardPage() { const queryClient = useQueryClient(); const { user } = useAuth(); const [isCreateDialogOpen, setIsCreateDialogOpen] = useState(false); - const [loadedSpendKeys, setLoadedSpendKeys] = useState>(new Set()); // Fetch private AI keys using React Query const { data: privateAIKeys = [] } = useQuery({ @@ -48,13 +47,8 @@ export default function DashboardPage() { refetchOnWindowFocus: false, // Prevent unnecessary refetches }); - // Use shared hook for data fetching - const { teamDetails, teamMembers, spendMap, regions } = usePrivateAIKeysData(privateAIKeys, loadedSpendKeys); - - // Load spend for a key - const loadSpend = useCallback(async (keyId: number) => { - setLoadedSpendKeys(prev => new Set([...prev, keyId])); - }, []); + // Use shared hook for data fetching (only for team details and regions) + const { teamDetails, teamMembers, regions } = usePrivateAIKeysData(privateAIKeys, new Set()); // Update budget period mutation const updateBudgetMutation = useMutation({ @@ -63,10 +57,8 @@ export default function DashboardPage() { return response.json(); }, onSuccess: (_, { keyId }) => { - // Refresh spend data - setLoadedSpendKeys(prev => new Set([...prev, keyId])); - // Invalidate spend queries - queryClient.invalidateQueries({ queryKey: ['private-ai-keys-spend'] }); + // Invalidate the specific key's spend query to refresh the data + queryClient.invalidateQueries({ queryKey: ['private-ai-key-spend', keyId] }); toast({ title: 'Success', description: 'Budget period updated successfully', @@ -174,8 +166,6 @@ export default function DashboardPage() { isLoading={createKeyMutation.isPending} showOwner={true} allowModification={false} - spendMap={spendMap} - onLoadSpend={loadSpend} onUpdateBudget={(keyId, budgetDuration) => updateBudgetMutation.mutate({ keyId, budgetDuration })} isDeleting={deleteKeyMutation.isPending} isUpdatingBudget={updateBudgetMutation.isPending} diff --git a/frontend/src/app/team-admin/private-ai-keys/page.tsx b/frontend/src/app/team-admin/private-ai-keys/page.tsx index e8f554f..b550dec 100644 --- a/frontend/src/app/team-admin/private-ai-keys/page.tsx +++ b/frontend/src/app/team-admin/private-ai-keys/page.tsx @@ -1,6 +1,6 @@ 'use client'; -import { useState } from 'react'; +import { useState, useCallback } from 'react'; import { useQuery, useMutation, useQueryClient } from '@tanstack/react-query'; import { useToast } from '@/hooks/use-toast'; import { get, post, del, put } from '@/utils/api'; @@ -31,39 +31,43 @@ interface TeamUser { export default function TeamAIKeysPage() { const { toast } = useToast(); + const queryClient = useQueryClient(); const { user } = useAuth(); const [isAddingKey, setIsAddingKey] = useState(false); - const [loadedSpendKeys, setLoadedSpendKeys] = useState>(new Set()); - - const queryClient = useQueryClient(); - const { data: keys = [], isLoading: isLoadingKeys } = useQuery({ - queryKey: ['private-ai-keys', user?.team_id], + // Fetch team members + const { data: teamMembersFull = [] } = useQuery({ + queryKey: ['team-members'], queryFn: async () => { - const response = await get(`private-ai-keys?team_id=${user?.team_id}`, { credentials: 'include' }); - const data = await response.json(); - return data; + const response = await get('teams/members', { credentials: 'include' }); + return response.json(); }, - enabled: !!user?.team_id, }); - // Use shared hook for data fetching - const { teamDetails, teamMembers, spendMap, regions } = usePrivateAIKeysData(keys, loadedSpendKeys); - - const { data: teamMembersFull = [] } = useQuery({ - queryKey: ['team-users'], + // Fetch private AI keys + const { data: keys = [], isLoading: isLoadingKeys } = useQuery({ + queryKey: ['private-ai-keys'], queryFn: async () => { - const response = await get('users', { credentials: 'include' }); - const allUsers = await response.json(); - return allUsers; + const response = await get('private-ai-keys', { credentials: 'include' }); + return response.json(); }, }); + // Use shared hook for data fetching (only for team details and regions) + const { teamDetails, teamMembers, regions } = usePrivateAIKeysData(keys, new Set()); + + // Create key mutation const createKeyMutation = useMutation({ - mutationFn: async (data: { name: string; region_id: number; owner_id?: number; team_id?: number; key_type: 'full' | 'llm' | 'vector' }) => { + mutationFn: async (data: { + name: string + region_id: number + key_type: 'full' | 'llm' | 'vector' + owner_id?: number + team_id?: number + }) => { const endpoint = data.key_type === 'full' ? 'private-ai-keys' : - data.key_type === 'llm' ? 'private-ai-keys/token' : - 'private-ai-keys/vector-db'; + data.key_type === 'llm' ? 'private-ai-keys/token' : + 'private-ai-keys/vector-db'; const response = await post(endpoint, data, { credentials: 'include' }); return response.json(); }, @@ -73,18 +77,19 @@ export default function TeamAIKeysPage() { setIsAddingKey(false); toast({ title: 'Success', - description: 'AI key added successfully', + description: 'AI key created successfully', }); }, - onError: () => { + onError: (error: Error) => { toast({ title: 'Error', - description: 'Failed to add AI key', + description: error.message, variant: 'destructive', }); }, }); + // Delete key mutation const deleteKeyMutation = useMutation({ mutationFn: async (keyId: number) => { const response = await del(`private-ai-keys/${keyId}`, { credentials: 'include' }); @@ -116,11 +121,8 @@ export default function TeamAIKeysPage() { return response.json(); }, onSuccess: (data, variables) => { - // Update the spend information for this specific key - queryClient.setQueryData(['private-ai-keys-spend', Array.from(loadedSpendKeys)], (oldData: Record = {}) => ({ - ...oldData, - [variables.keyId]: data - })); + // Invalidate the specific key's spend query to refresh the data + queryClient.invalidateQueries({ queryKey: ['private-ai-key-spend', variables.keyId] }); toast({ title: 'Success', description: 'Budget period updated successfully', @@ -175,8 +177,6 @@ export default function TeamAIKeysPage() { isDeleting={deleteKeyMutation.isPending} allowModification={true} showOwner={true} - spendMap={spendMap} - onLoadSpend={(keyId) => setLoadedSpendKeys(prev => new Set([...prev, keyId]))} onUpdateBudget={(keyId, budgetDuration) => { updateBudgetPeriodMutation.mutate({ keyId, budgetDuration }); }} diff --git a/frontend/src/components/private-ai-key-spend-cell.tsx b/frontend/src/components/private-ai-key-spend-cell.tsx new file mode 100644 index 0000000..8c1c115 --- /dev/null +++ b/frontend/src/components/private-ai-key-spend-cell.tsx @@ -0,0 +1,192 @@ +import { useState } from 'react'; +import { useQuery } from '@tanstack/react-query'; +import { Button } from '@/components/ui/button'; +import { Loader2, RefreshCw, Pencil } from 'lucide-react'; +import { formatTimeUntil } from '@/lib/utils'; +import { get } from '@/utils/api'; +import { + Dialog, + DialogContent, + DialogDescription, + DialogFooter, + DialogHeader, + DialogTitle, + DialogTrigger, +} from '@/components/ui/dialog'; +import { Input } from '@/components/ui/input'; +import { Label } from '@/components/ui/label'; + +interface SpendInfo { + spend: number; + expires: string; + created_at: string; + updated_at: string; + max_budget: number | null; + budget_duration: string | null; + budget_reset_at: string | null; +} + +interface PrivateAIKeySpendCellProps { + keyId: number; + hasLiteLLMToken: boolean; + allowModification?: boolean; + onUpdateBudget?: (keyId: number, budgetDuration: string) => void; + isUpdatingBudget?: boolean; +} + +export function PrivateAIKeySpendCell({ + keyId, + hasLiteLLMToken, + allowModification = false, + onUpdateBudget, + isUpdatingBudget = false, +}: PrivateAIKeySpendCellProps) { + const [isLoaded, setIsLoaded] = useState(false); + const [isRefreshing, setIsRefreshing] = useState(false); + const [openBudgetDialog, setOpenBudgetDialog] = useState(false); + + // Query for spend data - only enabled when isLoaded is true + const { data: spendData, isLoading, refetch } = useQuery({ + queryKey: ['private-ai-key-spend', keyId], + queryFn: async () => { + const response = await get(`private-ai-keys/${keyId}/spend`); + return response.json(); + }, + enabled: isLoaded, + }); + + const handleLoadSpend = () => { + setIsLoaded(true); + }; + + const handleRefreshSpend = async () => { + setIsRefreshing(true); + try { + await refetch(); + } catch (error) { + console.error('Failed to refresh spend data:', error); + } finally { + setIsRefreshing(false); + } + }; + + if (!hasLiteLLMToken) { + return null; + } + + if (!isLoaded) { + return ( + + ); + } + + if (isLoading) { + return ( +
+ + Loading spend... +
+ ); + } + + if (!spendData) { + return ( +
+ Failed to load spend data +
+ ); + } + + return ( +
+
+ + ${spendData.spend.toFixed(2)} + + + {spendData.max_budget !== null + ? `/ $${spendData.max_budget.toFixed(2)}` + : '(No budget)'} + + +
+ + {spendData.budget_duration || 'No budget period'} + {spendData.budget_reset_at && ` • Resets ${formatTimeUntil(spendData.budget_reset_at)}`} + {allowModification && onUpdateBudget && ( + + + + + + + Update Budget Period + + Set the budget period for this key. Examples: "30d" (30 days), "24h" (24 hours), "60m" (60 minutes) + + +
+
+ + +
+
+ + + + + + )} + +
+ ); +} diff --git a/frontend/src/components/private-ai-keys-table.tsx b/frontend/src/components/private-ai-keys-table.tsx index 706ce37..ac812b7 100644 --- a/frontend/src/components/private-ai-keys-table.tsx +++ b/frontend/src/components/private-ai-keys-table.tsx @@ -11,22 +11,10 @@ import { useTablePagination, } from '@/components/ui/table'; import { DeleteConfirmationDialog } from '@/components/ui/delete-confirmation-dialog'; -import { - Dialog, - DialogContent, - DialogDescription, - DialogFooter, - DialogHeader, - DialogTitle, - DialogTrigger, -} from '@/components/ui/dialog'; - -import { Input } from '@/components/ui/input'; -import { Label } from '@/components/ui/label'; -import { Eye, EyeOff, Pencil, Loader2, ArrowUp, ArrowDown, ArrowUpDown } from 'lucide-react'; -import { formatTimeUntil } from '@/lib/utils'; +import { Eye, EyeOff, ArrowUp, ArrowDown, ArrowUpDown } from 'lucide-react'; import { PrivateAIKey } from '@/types/private-ai-key'; import { TableFilters, FilterField } from '@/components/ui/table-filters'; +import { PrivateAIKeySpendCell } from '@/components/private-ai-key-spend-cell'; type SortField = 'name' | 'region' | 'owner' | null; type SortDirection = 'asc' | 'desc'; @@ -38,13 +26,6 @@ interface PrivateAIKeysTableProps { isLoading?: boolean; showOwner?: boolean; allowModification?: boolean; - spendMap?: Record; - onLoadSpend?: (keyId: number) => void; onUpdateBudget?: (keyId: number, budgetDuration: string) => void; isDeleting?: boolean; isUpdatingBudget?: boolean; @@ -58,8 +39,6 @@ export function PrivateAIKeysTable({ isLoading = false, showOwner = false, allowModification = false, - spendMap = {}, - onLoadSpend, onUpdateBudget, isDeleting = false, isUpdatingBudget = false, @@ -67,7 +46,6 @@ export function PrivateAIKeysTable({ teamMembers = [], }: PrivateAIKeysTableProps) { const [showPassword, setShowPassword] = useState>({}); - const [openBudgetDialog, setOpenBudgetDialog] = useState(null); const [sortField, setSortField] = useState(null); const [sortDirection, setSortDirection] = useState('asc'); const [keyTypeFilter, setKeyTypeFilter] = useState('all'); @@ -383,79 +361,13 @@ export function PrivateAIKeysTable({ )} - {spendMap[key.id] ? ( -
-
- - ${spendMap[key.id].spend.toFixed(2)} - - - {spendMap[key.id]?.max_budget !== null - ? `/ $${spendMap[key.id]?.max_budget?.toFixed(2)}` - : '(No budget)'} - -
- - {spendMap[key.id].budget_duration || 'No budget period'} - {spendMap[key.id].budget_reset_at && ` • Resets ${formatTimeUntil(spendMap[key.id].budget_reset_at as string)}`} - {allowModification && onUpdateBudget && ( - setOpenBudgetDialog(open ? key.id : null)}> - - - - - - Update Budget Period - - Set the budget period for this key. Examples: "30d" (30 days), "24h" (24 hours), "60m" (60 minutes) - - -
-
- - -
-
- - - - - - )} - -
- ) : key.litellm_token ? ( - - ) : null} + {allowModification && ( diff --git a/frontend/src/hooks/use-private-ai-keys-data.ts b/frontend/src/hooks/use-private-ai-keys-data.ts index 2bed9fb..a8cda84 100644 --- a/frontend/src/hooks/use-private-ai-keys-data.ts +++ b/frontend/src/hooks/use-private-ai-keys-data.ts @@ -62,26 +62,20 @@ export function usePrivateAIKeysData( }, }); - // Create individual queries for each key - // Use the keys array length to ensure stable hook count - const spendQueries = keys.map((key, index) => - useQuery({ - queryKey: ['private-ai-key-spend', key.id], - queryFn: async () => { - const response = await get(`private-ai-keys/${key.id}/spend`); - return response.json(); - }, - enabled: loadedSpendKeys.has(key.id), - }) - ); - - // Combine all spend data into a single map - const spendMap: Record = {}; - keys.forEach((key, index) => { - const query = spendQueries[index]; - if (query.data && loadedSpendKeys.has(key.id)) { - spendMap[key.id] = query.data; - } + // Query to get spend information for loaded keys + // Use a stable query key to avoid Rules of Hooks violations + const loadedSpendKeysArray = Array.from(loadedSpendKeys).sort(); + const { data: spendMap = {} } = useQuery>({ + queryKey: ['private-ai-keys-spend', loadedSpendKeysArray], + queryFn: async () => { + const spendPromises = loadedSpendKeysArray.map(async (keyId) => { + const response = await get(`private-ai-keys/${keyId}/spend`); + return [keyId, await response.json()] as [number, SpendInfo]; + }); + const spendResults = await Promise.all(spendPromises); + return Object.fromEntries(spendResults); + }, + enabled: loadedSpendKeysArray.length > 0, }); // Fetch regions From a010a93def0551d809a75fe238e88bfbebc1c41d Mon Sep 17 00:00:00 2001 From: Pippa H Date: Wed, 13 Aug 2025 14:00:39 +0200 Subject: [PATCH 4/4] Refactor imports in private AI keys pages to streamline code - Removed unused `useCallback` import from multiple private AI keys page components. - Improved code clarity by simplifying import statements across the affected files. --- frontend/src/app/admin/private-ai-keys/page.tsx | 2 +- frontend/src/app/private-ai-keys/page.tsx | 2 +- frontend/src/app/team-admin/private-ai-keys/page.tsx | 12 +----------- 3 files changed, 3 insertions(+), 13 deletions(-) diff --git a/frontend/src/app/admin/private-ai-keys/page.tsx b/frontend/src/app/admin/private-ai-keys/page.tsx index f7df07f..9d64539 100644 --- a/frontend/src/app/admin/private-ai-keys/page.tsx +++ b/frontend/src/app/admin/private-ai-keys/page.tsx @@ -1,6 +1,6 @@ 'use client'; -import { useState, useCallback } from 'react'; +import { useState } from 'react'; import { useQuery, useMutation, useQueryClient } from '@tanstack/react-query'; import { Button } from '@/components/ui/button'; import { useToast } from '@/hooks/use-toast'; diff --git a/frontend/src/app/private-ai-keys/page.tsx b/frontend/src/app/private-ai-keys/page.tsx index b5e5d27..c8fe9ba 100644 --- a/frontend/src/app/private-ai-keys/page.tsx +++ b/frontend/src/app/private-ai-keys/page.tsx @@ -1,6 +1,6 @@ 'use client'; -import { useState, useCallback } from 'react'; +import { useState } from 'react'; import { useMutation, useQueryClient, useQuery } from '@tanstack/react-query'; import { Card, CardContent } from '@/components/ui/card'; import { useToast } from '@/hooks/use-toast'; diff --git a/frontend/src/app/team-admin/private-ai-keys/page.tsx b/frontend/src/app/team-admin/private-ai-keys/page.tsx index b550dec..dc79e5a 100644 --- a/frontend/src/app/team-admin/private-ai-keys/page.tsx +++ b/frontend/src/app/team-admin/private-ai-keys/page.tsx @@ -1,6 +1,6 @@ 'use client'; -import { useState, useCallback } from 'react'; +import { useState } from 'react'; import { useQuery, useMutation, useQueryClient } from '@tanstack/react-query'; import { useToast } from '@/hooks/use-toast'; import { get, post, del, put } from '@/utils/api'; @@ -10,16 +10,6 @@ import { CreateAIKeyDialog } from '@/components/create-ai-key-dialog'; import { PrivateAIKey } from '@/types/private-ai-key'; import { usePrivateAIKeysData } from '@/hooks/use-private-ai-keys-data'; -interface SpendInfo { - spend: number; - expires: string; - created_at: string; - updated_at: string; - max_budget: number | null; - budget_duration: string | null; - budget_reset_at: string | null; -} - interface TeamUser { id: number; email: string;