Simplify resource limits

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.

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")

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)

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<Team[]>({
     queryKey: ['teams'],
@@ -907,7 +912,7 @@ export default function TeamsPage() {
                             >
                               {team.is_active ? 'Active' : 'Inactive'}
                             </span>
-                            {isTeamExpired(team) && (
+                            {isTeamExpired(team) && !hasActiveProducts(team) && (
                               <span className="inline-flex items-center rounded-full px-2.5 py-0.5 text-xs font-medium bg-red-600 text-white">
                                 Expired
                               </span>
@@ -968,14 +973,14 @@ export default function TeamsPage() {
                                                   {expandedTeam.is_active ? "Active" : "Inactive"}
                                                 </Badge>
                                               </div>
-                                              {isTeamExpired(expandedTeam) && (
+                                              {isTeamExpired(expandedTeam) && !hasActiveProducts(expandedTeam) && (
                                                 <div>
                                                   <p className="text-sm font-medium text-muted-foreground">Expiration Status</p>
                                                   <Badge variant="destructive" className="bg-red-600 hover:bg-red-700">
                                                     Expired
                                                   </Badge>
                                                 </div>
-                                                                                            )}
+                                              )}
                                               {expandedTeam.is_always_free && (
                                                 <div>
                                                   <p className="text-sm font-medium text-muted-foreground">Always Free Status</p>

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
 

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