diff --git a/app/core/resource_limits.py b/app/core/resource_limits.py index 278c2f9..4d6fae8 100644 --- a/app/core/resource_limits.py +++ b/app/core/resource_limits.py @@ -185,14 +185,7 @@ def get_token_restrictions(db: Session, team_id: int) -> tuple[int, float, int]: logger.error(f"Team not found for team_id: {team_id}") raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Team not found") - if result.last_payment is None: - days_left_in_period = result.max_key_duration - else: - days_left_in_period = result.max_key_duration - ( - datetime.now(UTC) - max(result.created_at.replace(tzinfo=UTC), result.last_payment.replace(tzinfo=UTC)) - ).days - - return days_left_in_period, result.max_max_spend, result.max_rpm_limit + return result.max_key_duration, result.max_max_spend, result.max_rpm_limit def get_team_limits(db: Session, team_id: int): # TODO: Go through all products, and create a master list of the limits on all fields for this team. diff --git a/app/core/worker.py b/app/core/worker.py index 22cb057..c6f6f0d 100644 --- a/app/core/worker.py +++ b/app/core/worker.py @@ -368,7 +368,14 @@ async def monitor_team_keys( # Determine what the new budget_duration will be for logging new_budget_duration = data.get("budget_duration", current_budget_duration) logger.info(f"Key {key.id} budget update triggered: changing from {current_budget_duration}, {current_max_budget} to {new_budget_duration}, {max_budget_amount}") - await litellm_service.update_budget(**data) + + # Extract litellm_token and other parameters for update_budget call + litellm_token = data.pop("litellm_token") + budget_duration = data.get("budget_duration") + budget_amount = data.get("budget_amount") + + # Call update_budget with correct parameter order + await litellm_service.update_budget(litellm_token, budget_duration, budget_amount=budget_amount) logger.info(f"Updated key {key.id} budget settings") else: logger.info(f"Key {key.id} budget settings already match the expected values, no update needed") diff --git a/app/services/litellm.py b/app/services/litellm.py index 93ceb80..765d06e 100644 --- a/app/services/litellm.py +++ b/app/services/litellm.py @@ -52,7 +52,7 @@ async def create_key(self, email: str, name: str, user_id: int, team_id: str, du if user_id is not None: request_data["user_id"] = str(user_id) - logger.info("Making request to LiteLLM API to generate key") + logger.info(f"Making request to LiteLLM API to generate key with data: {request_data}") response = requests.post( f"{self.api_url}/key/generate", json=request_data, @@ -123,6 +123,7 @@ async def get_key_info(self, litellm_token: str) -> dict: } ) response.raise_for_status() + logger.info(f"LiteLLM key information: {response.json()}") return response.json() except requests.exceptions.RequestException as e: error_msg = str(e) diff --git a/frontend/src/app/admin/teams/page.tsx b/frontend/src/app/admin/teams/page.tsx index d38566d..a69fbae 100644 --- a/frontend/src/app/admin/teams/page.tsx +++ b/frontend/src/app/admin/teams/page.tsx @@ -176,6 +176,11 @@ export default function TeamsPage() { return createdAt < thirtyDaysAgo && lastPayment < thirtyDaysAgo; }; + // Helper function to determine if a team has active products + const hasActiveProducts = (team: Team): boolean => { + return Boolean(team.products && team.products.some(product => product.active)); + }; + // Queries const { data: teams = [], isLoading: isLoadingTeams } = useQuery({ queryKey: ['teams'], @@ -907,7 +912,7 @@ export default function TeamsPage() { > {team.is_active ? 'Active' : 'Inactive'} - {isTeamExpired(team) && ( + {isTeamExpired(team) && !hasActiveProducts(team) && ( Expired @@ -968,14 +973,14 @@ export default function TeamsPage() { {expandedTeam.is_active ? "Active" : "Inactive"} - {isTeamExpired(expandedTeam) && ( + {isTeamExpired(expandedTeam) && !hasActiveProducts(expandedTeam) && (

Expiration Status

Expired
- )} + )} {expandedTeam.is_always_free && (

Always Free Status

diff --git a/tests/test_resource_limits.py b/tests/test_resource_limits.py index a72c360..72e0524 100644 --- a/tests/test_resource_limits.py +++ b/tests/test_resource_limits.py @@ -980,8 +980,8 @@ def test_get_token_restrictions_with_payment_history(db, test_team, test_product days_left, max_spend, rpm_limit = get_token_restrictions(db, test_team.id) - # Should have 15 days left (30 - 15) - assert days_left == 15 + # Should return the product's renewal_period_days, not calculated days left + assert days_left == test_product.renewal_period_days # 30 days assert max_spend == test_product.max_budget_per_key assert rpm_limit == test_product.rpm_per_key diff --git a/tests/test_worker.py b/tests/test_worker.py index 8767ff1..19d7f1a 100644 --- a/tests/test_worker.py +++ b/tests/test_worker.py @@ -987,13 +987,13 @@ async def test_monitor_team_keys_with_renewal_period_updates(mock_litellm, db, t # Check the first call (team key) first_call = mock_instance.update_budget.call_args_list[0] - assert first_call[1]['litellm_token'] == "team_token_123" + assert first_call[0][0] == "team_token_123" # First positional argument should be litellm_token assert first_call[1]['budget_amount'] == test_product.max_budget_per_key # budget_duration should not be updated since it's not None and reset time doesn't align # Check the second call (user key) second_call = mock_instance.update_budget.call_args_list[1] - assert second_call[1]['litellm_token'] == "user_token_456" + assert second_call[0][0] == "user_token_456" # First positional argument should be litellm_token assert second_call[1]['budget_amount'] == test_product.max_budget_per_key # budget_duration should not be updated since it's not None and reset time doesn't align @@ -1077,22 +1077,22 @@ async def test_monitor_team_keys_with_renewal_period_updates_no_products(mock_li # Verify get_key_info was called for both keys assert mock_instance.get_key_info.call_count == 2 - # Verify update_budget was called for both keys since they have different settings + # Verify update_budget was called for both keys since they have None budget_duration assert mock_instance.update_budget.call_count == 2 # Check the first call (team key) first_call = mock_instance.update_budget.call_args_list[0] - assert first_call[1]['litellm_token'] == "team_token_123" - assert first_call[1]['budget_duration'] == "30d" + assert first_call[0][0] == "team_token_123" # First positional argument should be litellm_token + assert first_call[0][1] == "30d" # Second positional argument should be budget_duration # Should not have budget_amount since no products were found - assert 'budget_amount' not in first_call[1] + assert first_call[1]['budget_amount'] is None # Check the second call (user key) second_call = mock_instance.update_budget.call_args_list[1] - assert second_call[1]['litellm_token'] == "user_token_456" - assert second_call[1]['budget_duration'] == "30d" + assert second_call[0][0] == "user_token_456" # First positional argument should be litellm_token + assert second_call[0][1] == "30d" # Second positional argument should be budget_duration # Should not have budget_amount since no products were found - assert 'budget_amount' not in second_call[1] + assert second_call[1]['budget_amount'] is None # Verify team total spend is calculated correctly assert team_total == 5.0 # 0.0 + 5.0 @@ -1156,8 +1156,8 @@ async def test_monitor_team_keys_reset_time_alignment_heuristic(mock_litellm, db # 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[0][0] == "team_token_123" # First positional argument should be litellm_token + assert update_call[0][1] == f"{test_product.renewal_period_days}d" # Second positional argument should be budget_duration assert update_call[1]['budget_amount'] == test_product.max_budget_per_key # Verify team total spend is calculated correctly @@ -1225,8 +1225,8 @@ async def test_monitor_team_keys_none_budget_duration_handled(mock_litellm, db, # 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[0][0] == "team_token_123" # First positional argument should be litellm_token + assert update_call[0][1] == f"{test_product.renewal_period_days}d" # Second positional argument should be budget_duration assert update_call[1]['budget_amount'] == test_product.max_budget_per_key # Verify team total spend is calculated correctly @@ -1353,9 +1353,9 @@ async def test_monitor_team_keys_both_mismatch_with_reset_time_alignment(mock_li # Verify update_budget was called with both budget_amount and budget_duration 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[0][0] == "team_token_123" # First positional argument should be litellm_token + assert update_call[0][1] == f"{test_product.renewal_period_days}d" # Second positional argument should be budget_duration assert update_call[1]['budget_amount'] == test_product.max_budget_per_key - assert update_call[1]['budget_duration'] == f"{test_product.renewal_period_days}d" # Verify team total spend is calculated correctly assert team_total == 10.0 @@ -1418,10 +1418,9 @@ async def test_monitor_team_keys_both_mismatch_without_reset_time_alignment(mock # Verify update_budget was called with only budget_amount 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[0][0] == "team_token_123" # First positional argument should be litellm_token + assert update_call[0][1] is None # Second positional argument should be None (no budget_duration update) assert update_call[1]['budget_amount'] == test_product.max_budget_per_key - # Should not include budget_duration since reset time doesn't align - assert 'budget_duration' not in update_call[1] # Verify team total spend is calculated correctly assert team_total == 10.0 @@ -1501,8 +1500,8 @@ async def test_monitor_team_keys_duration_parsing_different_units(mock_litellm, if should_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[0][0] == "team_token_123" # First positional argument should be litellm_token + assert update_call[0][1] == f"{test_product.renewal_period_days}d" # Second positional argument should be budget_duration else: assert mock_instance.update_budget.call_count == 0 @@ -1567,9 +1566,67 @@ async def test_monitor_team_keys_zero_duration_renewal(mock_litellm, db, test_te # Verify update_budget was called to fix the "0d" duration 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[0][0] == "team_token_123" # First positional argument should be litellm_token + assert update_call[0][1] == f"{test_product.renewal_period_days}d" # Second positional argument should be budget_duration 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_update_budget_parameter_issue(mock_litellm, db, test_team, test_product, test_region, test_team_user, test_team_key_creator): + """ + Test that update_budget is called with correct parameters when budget amount needs updating. + + GIVEN: A team with a product and keys that have different budget amounts + WHEN: monitor_team_keys is called with renewal period and budget amount + THEN: update_budget should be called with litellm_token as first positional argument, not as keyword argument + """ + # Set up team-product association + 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() + + # Mock LiteLLM service + mock_instance = mock_litellm.return_value + mock_instance.get_key_info = AsyncMock() + mock_instance.update_budget = AsyncMock() + + # Mock key info response - different budget amount triggers update + mock_instance.get_key_info.return_value = { + "info": { + "budget_reset_at": (datetime.now(UTC) + timedelta(days=30)).isoformat(), + "key_alias": "test_key", + "spend": 0.0, + "max_budget": 27.0, # Different from expected (120.0) + "budget_duration": "30d" + } + } + + # Get keys by region + keys_by_region = get_team_keys_by_region(db, test_team.id) + + # Call the 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 update_budget was called with correct parameters + assert mock_instance.update_budget.call_count == 1 + + # Check that litellm_token is passed as first positional argument, not as keyword + call_args = mock_instance.update_budget.call_args + # After the fix, litellm_token should be the first positional argument + assert call_args[0][0] == "team_token_123" # First positional argument should be litellm_token + assert call_args[1]['budget_amount'] == test_product.max_budget_per_key # budget_amount as keyword argument