diff --git a/api/controllers/console/billing/billing.py b/api/controllers/console/billing/billing.py index 4b0c82ae6c..cfe2c164c1 100644 --- a/api/controllers/console/billing/billing.py +++ b/api/controllers/console/billing/billing.py @@ -2,9 +2,12 @@ from flask_login import current_user from flask_restful import Resource, reqparse from controllers.console import api +from controllers.console.app.error import AppUnavailableError from controllers.console.wraps import account_initialization_required, only_edition_cloud, setup_required +from libs.helper import return_app_entity from libs.login import login_required from services.billing_service import BillingService +from services.quota_service import QuotaService class Subscription(Resource): @@ -37,3 +40,21 @@ class Invoices(Resource): api.add_resource(Subscription, "/billing/subscription") api.add_resource(Invoices, "/billing/invoices") + + +class QuotaUsage(Resource): + @setup_required + @login_required + @account_initialization_required + @only_edition_cloud + def get(self): + tenant = current_user.current_tenant + if not tenant: + raise AppUnavailableError() + + BillingService.is_tenant_owner_or_admin(current_user) + + quota_summary = QuotaService.get_quota_usage_summary(tenant) + return return_app_entity(quota_summary) + +api.add_resource(QuotaUsage, '/billing/quota/usage') diff --git a/api/models/account.py b/api/models/account.py index 7ffeefa980..9d4ed0e291 100644 --- a/api/models/account.py +++ b/api/models/account.py @@ -4,7 +4,7 @@ from typing import Optional, cast from flask_login import UserMixin # type: ignore from sqlalchemy import func -from sqlalchemy.orm import Mapped, mapped_column, reconstructor +from sqlalchemy.orm import Mapped, mapped_column, reconstructor, relationship from models.base import Base @@ -204,6 +204,8 @@ class Tenant(Base): custom_config = db.Column(db.Text) created_at = db.Column(db.DateTime, nullable=False, server_default=func.current_timestamp()) updated_at = db.Column(db.DateTime, nullable=False, server_default=func.current_timestamp()) + tenant_quota = relationship("TenantQuota", uselist=False, backref="tenant") + _quota_info: Optional[dict] = None def get_accounts(self) -> list[Account]: return ( @@ -220,6 +222,13 @@ class Tenant(Base): def custom_config_dict(self, value: dict): self.custom_config = json.dumps(value) + @property + def quota(self) -> dict: + if self._quota_info is None: + from services.billing_service import BillingService + self._quota_info = BillingService.get_tenant_quota(self.id) + return self._quota_info + class TenantAccountJoin(Base): __tablename__ = "tenant_account_joins" @@ -299,3 +308,23 @@ class TenantPluginPermission(Base): db.String(16), nullable=False, server_default="everyone" ) debug_permission: Mapped[DebugPermission] = mapped_column(db.String(16), nullable=False, server_default="noone") + + +class TenantQuota(Base): + __tablename__ = 'tenant_quotas' + __table_args__ = ( + db.PrimaryKeyConstraint('id', name='tenant_quota_pkey'), + db.Index('tenant_quota_tenant_id_idx', 'tenant_id'), + ) + + id: Mapped[str] = mapped_column(StringUUID, server_default=db.text("uuid_generate_v4()")) + tenant_id: Mapped[str] = mapped_column(StringUUID, db.ForeignKey('tenants.id'), unique=True, nullable=False) + max_users: Mapped[int] = mapped_column(db.Integer, nullable=True) + max_documents: Mapped[int] = mapped_column(db.Integer, nullable=True) + max_document_size_mb: Mapped[int] = mapped_column(db.Integer, nullable=True) + max_api_calls_per_day: Mapped[int] = mapped_column(db.Integer, nullable=True) + max_api_calls_per_month: Mapped[int] = mapped_column(db.Integer, nullable=True) + max_apps: Mapped[int] = mapped_column(db.Integer, nullable=True) + max_datasets: Mapped[int] = mapped_column(db.Integer, nullable=True) + created_at: Mapped[db.DateTime] = mapped_column(db.DateTime, nullable=False, server_default=func.current_timestamp()) + updated_at: Mapped[db.DateTime] = mapped_column(db.DateTime, nullable=False, server_default=func.current_timestamp(), onupdate=func.current_timestamp()) diff --git a/api/services/account_service.py b/api/services/account_service.py index ac84a46299..d3733c1a0a 100644 --- a/api/services/account_service.py +++ b/api/services/account_service.py @@ -15,6 +15,8 @@ from werkzeug.exceptions import Unauthorized from configs import dify_config from constants.languages import language_timezone_mapping, languages +from services.quota_service import QuotaService +from services.errors.quota import QuotaExceededError from events.tenant_event import tenant_was_created from extensions.ext_database import db from extensions.ext_redis import redis_client @@ -644,6 +646,12 @@ class TenantService: @staticmethod def create_tenant_member(tenant: Tenant, account: Account, role: str = "normal") -> TenantAccountJoin: """Create tenant member""" + + # Check user quota before adding a new member + current_user_count = db.session.query(TenantAccountJoin).filter(TenantAccountJoin.tenant_id == tenant.id).count() + if QuotaService.check_quota(tenant, 'max_users', current_user_count): + QuotaService.handle_quota_overage(tenant, 'max_users') + if role == TenantAccountRole.OWNER.value: if TenantService.has_roles(tenant, [TenantAccountRole.OWNER]): logging.error(f"Tenant {tenant.id} has already an owner.") diff --git a/api/services/app_service.py b/api/services/app_service.py index ebebf8fa58..512bd905d4 100644 --- a/api/services/app_service.py +++ b/api/services/app_service.py @@ -8,6 +8,8 @@ from flask_sqlalchemy.pagination import Pagination from configs import dify_config from constants.model_template import default_app_templates +from services.quota_service import QuotaService +from services.errors.quota import QuotaExceededError from core.agent.entities import AgentToolEntity from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError from core.model_manager import ModelManager @@ -78,6 +80,17 @@ class AppService: :param args: request args :param account: Account instance """ + # Check app quota + tenant = account.current_tenant + if tenant: + current_app_count = db.session.query(App).filter(App.tenant_id == tenant.id).count() + if QuotaService.check_quota(tenant, 'max_apps', current_app_count): + QuotaService.handle_quota_overage(tenant, 'max_apps') + else: + # Handle case where tenant is not found, though this should ideally not happen + # for a logged-in user creating an app. + pass + app_mode = AppMode.value_of(args["mode"]) app_template = default_app_templates[app_mode] diff --git a/api/services/billing_service.py b/api/services/billing_service.py index d44483ad89..b6eb87bcc9 100644 --- a/api/services/billing_service.py +++ b/api/services/billing_service.py @@ -148,6 +148,17 @@ class BillingService: params = {"keywords": keywords, "page": page, "limit": limit} return BillingService._send_request("GET", "/education/autocomplete", params=params) + @classmethod + def get_tenant_quota(cls, tenant_id: str) -> dict: + # In a real application, this would make a request to a billing API + # For this example, we'll use a mock response + # Replace with actual API call if needed + if tenant_id == "test_tenant_id_with_quota": + return {"max_users": 10, "max_documents": 1000} + else: + # Default quota for other tenants or if the API call fails + return {"max_users": 1, "max_documents": 10} + @classmethod def get_compliance_download_link( cls, diff --git a/api/services/dataset_service.py b/api/services/dataset_service.py index 6957315409..fc0fa247c0 100644 --- a/api/services/dataset_service.py +++ b/api/services/dataset_service.py @@ -14,6 +14,8 @@ from sqlalchemy.orm import Session from werkzeug.exceptions import NotFound from configs import dify_config +from services.quota_service import QuotaService +from services.errors.quota import QuotaExceededError from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError from core.model_manager import ModelManager from core.model_runtime.entities.model_entities import ModelType @@ -176,6 +178,17 @@ class DatasetService: embedding_model_name: Optional[str] = None, retrieval_model: Optional[RetrievalModel] = None, ): + # Check dataset quota + tenant = db.session.query(Tenant).filter(Tenant.id == tenant_id).first() + if tenant: + current_dataset_count = db.session.query(Dataset).filter(Dataset.tenant_id == tenant_id).count() + if QuotaService.check_quota(tenant, 'max_datasets', current_dataset_count): + QuotaService.handle_quota_overage(tenant, 'max_datasets') + else: + # Handle case where tenant is not found. This might indicate an issue. + # For now, let it proceed, or consider raising an error if tenant must exist. + pass + # check if dataset name already exists if db.session.query(Dataset).filter_by(name=name, tenant_id=tenant_id).first(): raise DatasetNameDuplicateError(f"Dataset with name {name} already exists.") @@ -923,7 +936,19 @@ class DocumentService: if count > batch_upload_limit: raise ValueError(f"You have reached the batch upload limit of {batch_upload_limit}.") - DocumentService.check_documents_upload_quota(count, features) + total_size_bytes = 0 + if knowledge_config.data_source.info_list.data_source_type == "upload_file": + upload_file_ids = knowledge_config.data_source.info_list.file_info_list.file_ids + for file_id in upload_file_ids: + upload_file = db.session.query(UploadFile).filter(UploadFile.id == file_id, UploadFile.tenant_id == dataset.tenant_id).first() + if upload_file: + total_size_bytes += upload_file.size + # For Notion and Website crawls, size is not directly available here. + # Passing 0 for size check, meaning this specific check will primarily apply to file uploads. + # Alternative: estimate size or skip check for these types if appropriate. + + total_size_mb = total_size_bytes / (1024 * 1024) + DocumentService.check_documents_upload_quota(total_size_mb, features) # if dataset is empty, update dataset data_source_type if not dataset.data_source_type: @@ -1180,12 +1205,19 @@ class DocumentService: return documents, batch @staticmethod - def check_documents_upload_quota(count: int, features: FeatureModel): - can_upload_size = features.documents_upload_quota.limit - features.documents_upload_quota.size - if count > can_upload_size: - raise ValueError( - f"You have reached the limit of your subscription. Only {can_upload_size} documents can be uploaded." - ) + def check_documents_upload_quota(size_of_current_upload_mb: int, features: FeatureModel): # Renamed 'count' to 'size_of_current_upload_mb' for clarity + """ + Checks if the tenant has exceeded its document size quota (max_document_size_mb) + for the current upload batch. + Note: The `features` parameter is currently unused after the change, + but kept for signature consistency if other parts of the system rely on it. + """ + tenant = current_user.current_tenant + + # Check for max_document_size_mb using QuotaService. + # This replaces the previous ValueError logic with a check against tenant quota. + if QuotaService.check_quota(tenant, 'max_document_size_mb', size_of_current_upload_mb): + QuotaService.handle_quota_overage(tenant, 'max_document_size_mb') @staticmethod def build_document( @@ -2250,6 +2282,11 @@ class SegmentService: if not dataset: raise NotFound("Dataset not found.") + # Check document quota + current_document_count = db.session.query(Document).filter(Document.dataset_id == dataset_id).count() + if QuotaService.check_quota(dataset.tenant, 'max_documents', current_document_count): + QuotaService.handle_quota_overage(dataset.tenant, 'max_documents') + # check user's model setting DatasetService.check_dataset_model_setting(dataset) diff --git a/api/services/errors/quota.py b/api/services/errors/quota.py new file mode 100644 index 0000000000..8643b8ea98 --- /dev/null +++ b/api/services/errors/quota.py @@ -0,0 +1,6 @@ +class QuotaExceededError(Exception): + """Custom exception for quota exceeded errors.""" + def __init__(self, resource: str, message: str = None): + self.resource = resource + self.message = message or f"Quota for resource '{resource}' exceeded." + super().__init__(self.message) diff --git a/api/services/quota_service.py b/api/services/quota_service.py new file mode 100644 index 0000000000..1b98503ae7 --- /dev/null +++ b/api/services/quota_service.py @@ -0,0 +1,71 @@ +from extensions.ext_database import db +from models.account import Tenant, TenantAccountJoin +from models.model import App +from models.dataset import Dataset, Document +from services.errors.quota import QuotaExceededError + +class QuotaService: + @staticmethod + def get_tenant_quota(tenant: Tenant) -> dict: + """ + Retrieves the tenant's quota information. + """ + return tenant.quota + + @staticmethod + def check_quota(tenant: Tenant, resource: str, current_usage: int) -> bool: + """ + Checks if the tenant has exceeded the quota for the given resource. + Returns True if quota is exceeded, False otherwise. + """ + quota_info = tenant.quota + resource_limit = quota_info.get(resource) + + if resource_limit is not None and current_usage >= resource_limit: + return True # Quota exceeded + return False # Quota not exceeded + + @staticmethod + def handle_quota_overage(tenant: Tenant, resource: str): + """ + Handles quota overages by raising a QuotaExceededError. + """ + # Additional logic for handling quota overage can be added here in the future + # For example, sending notifications, logging, etc. + raise QuotaExceededError(resource=resource, message=f"Quota for resource '{resource}' exceeded for tenant {tenant.id}.") + + @staticmethod + def get_quota_usage_summary(tenant: Tenant) -> dict: + """ + Retrieves the tenant's quota limits and current usage for various resources. + """ + quota_limits = tenant.quota + + # Calculate current usage + current_users = db.session.query(TenantAccountJoin).filter(TenantAccountJoin.tenant_id == tenant.id).count() + # Note: Document model might need to be specific to the tenant if not already filtered by dataset.tenant_id + # Assuming Document table has a direct tenant_id or can be joined through Dataset. + # For simplicity, let's assume Documents are implicitly tenant-specific via their Datasets. + current_documents = db.session.query(Document).join(Dataset, Document.dataset_id == Dataset.id).filter(Dataset.tenant_id == tenant.id).count() + current_apps = db.session.query(App).filter(App.tenant_id == tenant.id).count() + current_datasets = db.session.query(Dataset).filter(Dataset.tenant_id == tenant.id).count() + + # max_document_size_mb and api calls are not directly counted here as they are usage metrics, + # not simple counts of db entities. This summary focuses on countable entities. + # These could be added if specific tracking mechanisms for them exist. + + usage = { + "max_users": current_users, + "max_documents": current_documents, + "max_apps": current_apps, + "max_datasets": current_datasets, + # Placeholder for other usage metrics if they were to be counted + "max_document_size_mb": "N/A (refer to individual document sizes)", # Or a sum if meaningful + "max_api_calls_per_day": "N/A (tracked externally)", + "max_api_calls_per_month": "N/A (tracked externally)", + } + + return { + "limits": quota_limits, + "usage": usage + } diff --git a/api/tests/integration_tests/controllers/test_controllers.py b/api/tests/integration_tests/controllers/test_controllers.py index 276ad3a7ed..47b9888853 100644 --- a/api/tests/integration_tests/controllers/test_controllers.py +++ b/api/tests/integration_tests/controllers/test_controllers.py @@ -1,9 +1,86 @@ -from unittest.mock import patch +import json +from unittest.mock import patch, MagicMock -from app_fixture import mock_user # type: ignore +import pytest +# Assuming app_fixture.py is in the same directory or accessible via PYTHONPATH +# If app_fixture.py defines user_owner, user_normal etc., use them. +# For now, we'll create MagicMock users as needed within tests or use a generic mock_user. +from .app_fixture import mock_user as generic_mock_user_from_fixture +# It's better to have specific mock users for roles if app_fixture provides them. +# e.g. from .app_fixture import mock_owner_user, mock_normal_user + +from models.account import Tenant, TenantAccountRole +# QuotaService will be mocked at the service layer, not imported directly unless needed for type hinting +# from services.quota_service import QuotaService + + +# A basic mock user if specific ones are not available from app_fixture +def create_mock_user(role=TenantAccountRole.OWNER, tenant_id="test_tenant_id"): + user = MagicMock() + user.id = "test_user_id" + user.current_tenant_id = tenant_id + user.current_tenant = MagicMock(spec=Tenant) + user.current_tenant.id = tenant_id + # This role is on the TenantAccountJoin, not directly on Account. + # The BillingService.is_tenant_owner_or_admin will query TenantAccountJoin. + # So, for these tests, we might need to also patch BillingService.is_tenant_owner_or_admin + # or ensure the test setup correctly creates these DB entries if not fully mocking. + # For simplicity in this example, is_tenant_owner_or_admin will be patched. + return user + +mock_owner_user = create_mock_user(role=TenantAccountRole.OWNER) +mock_normal_user = create_mock_user(role=TenantAccountRole.NORMAL) + + +class TestBillingController: + + @patch('services.billing_service.BillingService.is_tenant_owner_or_admin', return_value=True) + @patch('services.quota_service.QuotaService.get_quota_usage_summary') + @patch('flask_login.utils._get_user') + def test_get_quota_usage_summary_owner_success(self, mock_get_user, mock_get_quota_summary, mock_is_admin, app): + mock_get_user.return_value = mock_owner_user # Simulate owner logged in + + sample_summary = { + 'limits': {'max_users': 10, 'max_apps': 5, 'max_datasets': 3, 'max_documents': 100, 'max_document_size_mb': 50, 'max_api_calls_per_day': 1000, 'max_api_calls_per_month': 10000}, + 'usage': {'max_users': 1, 'max_apps': 2, 'max_datasets': 1, 'max_documents': 10, 'max_document_size_mb': 'N/A', 'max_api_calls_per_day': 'N/A', 'max_api_calls_per_month': 'N/A'} + } + mock_get_quota_summary.return_value = sample_summary + + with app.test_client() as client: + response = client.get('/console/api/billing/quota/usage') -def test_post_requires_login(app): - with app.test_client() as client, patch("flask_login.utils._get_user", mock_user): - response = client.get("/console/api/data-source/integrates") assert response.status_code == 200 + response_data = json.loads(response.data.decode('utf-8')) + assert response_data['limits'] == sample_summary['limits'] + assert response_data['usage'] == sample_summary['usage'] + mock_get_quota_summary.assert_called_once_with(mock_owner_user.current_tenant) + mock_is_admin.assert_called_once_with(mock_owner_user) + + + def test_get_quota_usage_summary_unauthenticated(self, app): + with app.test_client() as client: + response = client.get('/console/api/billing/quota/usage') + # Should redirect to /login or return 401 if API + # Based on typical Flask-Login behavior for @login_required, + # it might redirect for browsers or return 401 for XHR. + # Let's assume it's configured to return 401 for API endpoints. + assert response.status_code == 401 + + + @patch('services.billing_service.BillingService.is_tenant_owner_or_admin', return_value=False) + @patch('flask_login.utils._get_user') + def test_get_quota_usage_summary_normal_user_forbidden(self, mock_get_user, mock_is_admin, app): + mock_get_user.return_value = mock_normal_user # Simulate normal user logged in + + with app.test_client() as client: + response = client.get('/console/api/billing/quota/usage') + + assert response.status_code == 403 + mock_is_admin.assert_called_once_with(mock_normal_user) + +# Keep existing test if it's still relevant, or remove/move if it belongs elsewhere +# def test_post_requires_login(app): +# with app.test_client() as client, patch("flask_login.utils._get_user", generic_mock_user_from_fixture): +# response = client.get("/console/api/data-source/integrates") +# assert response.status_code == 200 diff --git a/api/tests/unit_tests/services/test_quota_service.py b/api/tests/unit_tests/services/test_quota_service.py new file mode 100644 index 0000000000..8d06639f38 --- /dev/null +++ b/api/tests/unit_tests/services/test_quota_service.py @@ -0,0 +1,139 @@ +import unittest +from unittest.mock import patch, MagicMock, PropertyMock + +from extensions.ext_database import db # Needed for mocking db.session +from models.account import Tenant, TenantAccountJoin +from models.model import App +from models.dataset import Dataset, Document +from services.quota_service import QuotaService +from services.errors.quota import QuotaExceededError + +class TestQuotaService(unittest.TestCase): + + def test_get_tenant_quota(self): + mock_tenant = MagicMock(spec=Tenant) + mock_quota_data = {"max_users": 10, "max_apps": 5} + type(mock_tenant).quota = PropertyMock(return_value=mock_quota_data) + + result = QuotaService.get_tenant_quota(mock_tenant) + self.assertEqual(result, mock_quota_data) + + def test_check_quota(self): + mock_tenant = MagicMock(spec=Tenant) + + # Scenario 1: Quota not exceeded + type(mock_tenant).quota = PropertyMock(return_value={"max_users": 10}) + self.assertFalse(QuotaService.check_quota(mock_tenant, "max_users", 5)) + + # Scenario 2: Quota exceeded + type(mock_tenant).quota = PropertyMock(return_value={"max_users": 10}) + self.assertTrue(QuotaService.check_quota(mock_tenant, "max_users", 10)) + self.assertTrue(QuotaService.check_quota(mock_tenant, "max_users", 11)) + + # Scenario 3: Quota limit is None (unlimited) + type(mock_tenant).quota = PropertyMock(return_value={"max_users": None}) + self.assertFalse(QuotaService.check_quota(mock_tenant, "max_users", 100)) + + # Scenario 4: Unknown resource + type(mock_tenant).quota = PropertyMock(return_value={"max_users": 10}) + self.assertFalse(QuotaService.check_quota(mock_tenant, "unknown_resource", 5)) + + # Scenario 5: Resource limit is 0 + type(mock_tenant).quota = PropertyMock(return_value={"max_users": 0}) + self.assertTrue(QuotaService.check_quota(mock_tenant, "max_users", 0)) + self.assertTrue(QuotaService.check_quota(mock_tenant, "max_users", 1)) + + # Scenario 6: Resource not in quota (should be treated as no limit / not enforced) + type(mock_tenant).quota = PropertyMock(return_value={}) # empty dict + self.assertFalse(QuotaService.check_quota(mock_tenant, "max_users", 10)) + + + def test_handle_quota_overage(self): + mock_tenant = MagicMock(spec=Tenant) + mock_tenant.id = "test_tenant_id" + resource_name = "max_apps" + + with self.assertRaises(QuotaExceededError) as context: + QuotaService.handle_quota_overage(mock_tenant, resource_name) + + self.assertEqual(context.exception.resource, resource_name) + self.assertIn(resource_name, context.exception.message) + self.assertIn(mock_tenant.id, context.exception.message) + + @patch('extensions.ext_database.db.session.query') + def test_get_quota_usage_summary(self, mock_db_query): + mock_tenant = MagicMock(spec=Tenant) + mock_tenant.id = "tenant_123" + mock_limits = { + "max_users": 10, + "max_documents": 100, + "max_document_size_mb": 50, + "max_api_calls_per_day": 1000, + "max_api_calls_per_month": 10000, + "max_apps": 5, + "max_datasets": 3, + } + type(mock_tenant).quota = PropertyMock(return_value=mock_limits) + + # Mocking chain for db.session.query(...).filter(...).count() + # Users count + mock_query_users = MagicMock() + mock_query_users.filter.return_value.count.return_value = 7 + + # Documents count + mock_query_documents = MagicMock() + mock_query_documents.join.return_value.filter.return_value.count.return_value = 50 + + # Apps count + mock_query_apps = MagicMock() + mock_query_apps.filter.return_value.count.return_value = 3 + + # Datasets count + mock_query_datasets = MagicMock() + mock_query_datasets.filter.return_value.count.return_value = 2 + + # Configure mock_db_query to return the correct mock based on the model + def side_effect(model): + if model == TenantAccountJoin: + return mock_query_users + elif model == Document: + return mock_query_documents + elif model == App: + return mock_query_apps + elif model == Dataset: + return mock_query_datasets + return MagicMock() # Default mock for any other queries + + mock_db_query.side_effect = side_effect + + summary = QuotaService.get_quota_usage_summary(mock_tenant) + + expected_usage = { + "max_users": 7, + "max_documents": 50, + "max_apps": 3, + "max_datasets": 2, + "max_document_size_mb": "N/A (refer to individual document sizes)", + "max_api_calls_per_day": "N/A (tracked externally)", + "max_api_calls_per_month": "N/A (tracked externally)", + } + + self.assertEqual(summary["limits"], mock_limits) + self.assertEqual(summary["usage"], expected_usage) + + # Verify calls to db.session.query + mock_db_query.assert_any_call(TenantAccountJoin) + mock_query_users.filter.assert_called_with(TenantAccountJoin.tenant_id == mock_tenant.id) + + mock_db_query.assert_any_call(Document) + mock_query_documents.join.assert_called_with(Dataset, Document.dataset_id == Dataset.id) + mock_query_documents.join.return_value.filter.assert_called_with(Dataset.tenant_id == mock_tenant.id) + + mock_db_query.assert_any_call(App) + mock_query_apps.filter.assert_called_with(App.tenant_id == mock_tenant.id) + + mock_db_query.assert_any_call(Dataset) + mock_query_datasets.filter.assert_called_with(Dataset.tenant_id == mock_tenant.id) + +if __name__ == '__main__': + unittest.main() diff --git a/migrations/versions/20240424100000_create_tenant_quotas_table.py b/migrations/versions/20240424100000_create_tenant_quotas_table.py new file mode 100644 index 0000000000..734f99a703 --- /dev/null +++ b/migrations/versions/20240424100000_create_tenant_quotas_table.py @@ -0,0 +1,37 @@ +"""create_tenant_quotas_table + +Revision ID: 20240424100000 +Revises: +Create Date: 2024-04-24 10:00:00.000000 + +""" +from alembic import op +import sqlalchemy as sa +from sqlalchemy.dialects import postgresql + +# revision identifiers, used by Alembic. +revision = '20240424100000' +down_revision = None +branch_labels = None +depends_on = None + + +def upgrade(): + op.create_table( + 'tenant_quotas', + sa.Column('id', postgresql.UUID(as_uuid=True), primary_key=True, server_default=sa.text('gen_random_uuid()')), + sa.Column('tenant_id', postgresql.UUID(as_uuid=True), sa.ForeignKey('tenants.id'), unique=True, nullable=False), + sa.Column('max_users', sa.Integer(), nullable=True), + sa.Column('max_documents', sa.Integer(), nullable=True), + sa.Column('max_document_size_mb', sa.Integer(), nullable=True), + sa.Column('max_api_calls_per_day', sa.Integer(), nullable=True), + sa.Column('max_api_calls_per_month', sa.Integer(), nullable=True), + sa.Column('max_apps', sa.Integer(), nullable=True), + sa.Column('max_datasets', sa.Integer(), nullable=True), + sa.Column('created_at', sa.DateTime(), server_default=sa.func.current_timestamp(), nullable=False), + sa.Column('updated_at', sa.DateTime(), server_default=sa.func.current_timestamp(), onupdate=sa.func.current_timestamp(), nullable=False) + ) + + +def downgrade(): + op.drop_table('tenant_quotas') diff --git a/web/app/components/billing/plan/index.tsx b/web/app/components/billing/plan/index.tsx index 7badb3666f..985c5c4aeb 100644 --- a/web/app/components/billing/plan/index.tsx +++ b/web/app/components/billing/plan/index.tsx @@ -13,13 +13,14 @@ import { RiSquareLine, } from '@remixicon/react' import { Plan, SelfHostedPlan } from '../type' -import VectorSpaceInfo from '../usage-info/vector-space-info' -import AppsInfo from '../usage-info/apps-info' +// import VectorSpaceInfo from '../usage-info/vector-space-info' // Keep if still needed for other purposes, or remove if fully replaced +// import AppsInfo from '../usage-info/apps-info' // Keep if still needed for other purposes, or remove if fully replaced +import UsageInfoContainer from '../usage-info' // Import the new container import UpgradeBtn from '../upgrade-btn' import { useProviderContext } from '@/context/provider-context' import { useAppContext } from '@/context/app-context' import Button from '@/app/components/base/button' -import UsageInfo from '@/app/components/billing/usage-info' +// UsageInfo (single item) is no longer directly used here, it's used by UsageInfoContainer import VerifyStateModal from '@/app/education-apply/verify-state-modal' import { EDUCATION_VERIFYING_LOCALSTORAGE_ITEM } from '@/app/education-apply/constants' import { useEducationVerify } from '@/service/use-education' @@ -40,10 +41,11 @@ const PlanComp: FC = ({ type, } = plan - const { - usage, - total, - } = plan + // Remove usage and total from here as they are fetched by UsageInfoContainer + // const { + // usage, + // total, + // } = plan const [showModal, setShowModal] = React.useState(false) const { mutateAsync } = useEducationVerify() @@ -98,29 +100,9 @@ const PlanComp: FC = ({ - {/* Plan detail */} -
- - - - - - + {/* Plan detail - Replaced with UsageInfoContainer */} +
{/* Keep p-2 for padding, adjust grid if UsageInfoContainer handles its own grid */} +
{ + return get('/billing/quota/usage') +} + +// Add other billing related services here if any in the future +// For example: +// export const fetchSubscriptionPlans = () => { +// return get('/billing/plans') +// } +// +// export const updateSubscription = (planId: string) => { +// return post('/billing/subscription', { body: { plan_id: planId } }) +// } diff --git a/web/app/components/billing/type.ts b/web/app/components/billing/type.ts index 506ef46799..22de309cec 100644 --- a/web/app/components/billing/type.ts +++ b/web/app/components/billing/type.ts @@ -110,3 +110,28 @@ export type SubscriptionItem = { export type SubscriptionUrlsBackend = { url: string } + +export interface QuotaLimits { + max_users: number | null; + max_documents: number | null; + max_document_size_mb: number | null; + max_api_calls_per_day: number | null; + max_api_calls_per_month: number | null; + max_apps: number | null; + max_datasets: number | null; +} + +export interface CurrentUsage { + max_users: number; + max_documents: number; + max_document_size_mb: string; // Or number if it's a calculated sum + max_api_calls_per_day: string; // Or number if tracked + max_api_calls_per_month: string; // Or number if tracked + max_apps: number; + max_datasets: number; +} + +export interface QuotaUsageResponse { + limits: QuotaLimits; + usage: CurrentUsage; +} diff --git a/web/app/components/billing/usage-info/index.tsx b/web/app/components/billing/usage-info/index.tsx index 30b4bca776..e9bac141b5 100644 --- a/web/app/components/billing/usage-info/index.tsx +++ b/web/app/components/billing/usage-info/index.tsx @@ -1,26 +1,38 @@ 'use client' import type { FC } from 'react' -import React from 'react' +import React, { useEffect, useState, useCallback } from 'react' import { useTranslation } from 'react-i18next' import ProgressBar from '../progress-bar' import { NUM_INFINITE } from '../config' import Tooltip from '@/app/components/base/tooltip' import cn from '@/utils/classnames' +import { fetchQuotaUsage } from '../service' // Assuming service.ts is in the same directory +import type { QuotaUsageResponse, QuotaLimits, CurrentUsage } from '../type' +import { + UsersIcon, + DocumentTextIcon, + CubeIcon, + ServerStackIcon, + CpuChipIcon, + CloudArrowUpIcon, +} from '@heroicons/react/24/outline' +import Alert from '@/app/components/base/alert' -type Props = { +// Renamed original UsageInfo to UsageItem +type UsageItemProps = { className?: string - Icon: any + Icon: React.ElementType name: string tooltip?: string - usage: number - total: number + usage: number | string // Usage can be a string like "N/A" + total: number | null // Limit can be null for unlimited unit?: string } -const LOW = 50 -const MIDDLE = 80 +const LOW_USAGE_PERCENT = 50 +const MEDIUM_USAGE_PERCENT = 80 -const UsageInfo: FC = ({ +const UsageItem: FC = ({ className, Icon, name, @@ -31,16 +43,28 @@ const UsageInfo: FC = ({ }) => { const { t } = useTranslation() - const percent = usage / total * 100 + const isUnlimited = total === null || total === 0 || total === NUM_INFINITE + const numericUsage = typeof usage === 'number' ? usage : 0 + const numericTotal = typeof total === 'number' ? total : 0 + + let percent = 0 + if (!isUnlimited && numericTotal > 0 && typeof usage === 'number') + percent = (numericUsage / numericTotal) * 100 + else if (isUnlimited && typeof usage === 'number' && usage > 0) + percent = 0 // Show some minimal bar for used unlimited resources, or handle as 0 + + // If usage is "N/A" or similar, treat percent as 0 or handle as special case + if (typeof usage !== 'number') + percent = 0 + const color = (() => { - if (percent < LOW) + if (isUnlimited || percent < LOW_USAGE_PERCENT) return 'bg-components-progress-bar-progress-solid' - - if (percent < MIDDLE) + if (percent < MEDIUM_USAGE_PERCENT) return 'bg-components-progress-warning-progress' - return 'bg-components-progress-error-progress' })() + return (
@@ -48,18 +72,14 @@ const UsageInfo: FC = ({
{name}
{tooltip && ( - {tooltip} -
- } + popupContent={
{tooltip}
} /> )}
-
- {usage} +
+ {typeof usage === 'number' ? usage : t('billing.plansCommon.notAvailable')}
/
-
{total === NUM_INFINITE ? t('billing.plansCommon.unlimited') : `${total}${unit}`}
+
{isUnlimited ? t('billing.plansCommon.unlimited') : `${numericTotal}${unit}`}
= ({
) } -export default React.memo(UsageInfo) + +// New UsageInfo container component +const UsageInfoContainer: FC<{ className?: string }> = ({ className }) => { + const { t } = useTranslation() + const [quotaUsage, setQuotaUsage] = useState(null) + const [isLoading, setIsLoading] = useState(true) + const [error, setError] = useState(null) + + const loadQuotaUsage = useCallback(async () => { + setIsLoading(true) + setError(null) + try { + const data = await fetchQuotaUsage() + setQuotaUsage(data) + } + catch (e) { + setError(e) + console.error('Failed to fetch quota usage:', e) + } + finally { + setIsLoading(false) + } + }, []) + + useEffect(() => { + loadQuotaUsage() + }, [loadQuotaUsage]) + + // Resource key to display name and icon mapping + const resourceDisplayConfig: Record = { + max_users: { name: t('billing.usageInfo.teamMembers'), Icon: UsersIcon }, + max_documents: { name: t('billing.usageInfo.documents'), Icon: DocumentTextIcon }, + max_apps: { name: t('billing.usageInfo.apps'), Icon: CubeIcon }, + max_datasets: { name: t('billing.usageInfo.datasets'), Icon: ServerStackIcon }, + max_document_size_mb: { name: t('billing.usageInfo.documentProcessing'), Icon: CloudArrowUpIcon, unit: 'MB' }, // Assuming this is a limit per upload + max_api_calls_per_day: { name: t('billing.usageInfo.apiCallsPerDay'), Icon: CpuChipIcon }, + max_api_calls_per_month: { name: t('billing.usageInfo.apiCallsPerMonth'), Icon: CpuChipIcon }, + } + + if (isLoading) { + return ( +
+ {Array.from({ length: 4 }).map((_, i) => ( // Skeleton loaders +
+
+
+
+
+
+ ))} +
+ ) + } + + if (error) { + return ( + + {t('billing.usageInfo.fetchFailed')} {error.message || String(error)} + + ) + } + + if (!quotaUsage) + return null + + return ( +
+ {(Object.keys(quotaUsage.usage) as Array).map((key) => { + const config = resourceDisplayConfig[key] + if (!config) // Skip if no display config for this key + return null + + // Skip rendering items where usage is "N/A" and limit is also not set (or 0/null) + // This primarily targets API call limits which are marked "N/A" in current backend summary for usage + if (typeof quotaUsage.usage[key] !== 'number' && (!quotaUsage.limits[key] || quotaUsage.limits[key] === null)) { + if (key === 'max_document_size_mb' && typeof quotaUsage.limits[key] === 'number' && quotaUsage.limits[key]! > 0) { + // Special case for max_document_size_mb: show if limit exists, even if usage is "N/A" + } else { + return null; + } + } + + + return ( + + ) + })} +
+ ) +} + +export default React.memo(UsageInfoContainer)