mirror of
https://github.com/langgenius/dify.git
synced 2026-06-07 16:32:01 +08:00
chore: backend feature api exclude_vector_space (#36642)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
parent
5aa24c25d9
commit
f5d664887b
@ -77,7 +77,7 @@ register_response_schema_models(console_ns, SimpleResultDataResponse, Verificati
|
||||
def _is_role_enabled(role: TenantAccountRole | str, tenant_id: str) -> bool:
|
||||
if role != TenantAccountRole.DATASET_OPERATOR:
|
||||
return True
|
||||
return FeatureService.get_features(tenant_id=tenant_id).dataset_operator_enabled
|
||||
return FeatureService.get_features(tenant_id=tenant_id, exclude_vector_space=True).dataset_operator_enabled
|
||||
|
||||
|
||||
def _normalize_invitee_emails(emails: list[str]) -> list[str]:
|
||||
@ -113,7 +113,7 @@ def _check_member_invite_limits(tenant_id: str, new_member_count: int) -> None:
|
||||
if new_member_count <= 0:
|
||||
return
|
||||
|
||||
features = FeatureService.get_features(tenant_id=tenant_id)
|
||||
features = FeatureService.get_features(tenant_id=tenant_id, exclude_vector_space=True)
|
||||
|
||||
if dify_config.ENTERPRISE_ENABLED:
|
||||
workspace_members = features.workspace_members
|
||||
|
||||
@ -166,10 +166,10 @@ class TenantListApi(Resource):
|
||||
if tenant_plan:
|
||||
plan = tenant_plan["plan"] or CloudPlan.SANDBOX
|
||||
else:
|
||||
features = FeatureService.get_features(tenant.id)
|
||||
features = FeatureService.get_features(tenant.id, exclude_vector_space=True)
|
||||
plan = features.billing.subscription.plan or CloudPlan.SANDBOX
|
||||
elif not is_enterprise_only:
|
||||
features = FeatureService.get_features(tenant.id)
|
||||
features = FeatureService.get_features(tenant.id, exclude_vector_space=True)
|
||||
plan = features.billing.subscription.plan or CloudPlan.SANDBOX
|
||||
|
||||
# Create a dictionary with tenant attributes
|
||||
|
||||
@ -96,21 +96,28 @@ def cloud_edition_billing_resource_check[**P, R](resource: str) -> Callable[[Cal
|
||||
@wraps(view)
|
||||
def decorated(*args: P.args, **kwargs: P.kwargs):
|
||||
_, current_tenant_id = current_account_with_tenant()
|
||||
features = FeatureService.get_features(current_tenant_id)
|
||||
if resource == "vector_space":
|
||||
if not dify_config.BILLING_ENABLED:
|
||||
return view(*args, **kwargs)
|
||||
|
||||
vector_space = FeatureService.get_vector_space(current_tenant_id)
|
||||
if 0 < vector_space.limit <= vector_space.size:
|
||||
abort(
|
||||
403,
|
||||
"The capacity of the knowledge storage space has reached the limit of your subscription.",
|
||||
)
|
||||
return view(*args, **kwargs)
|
||||
|
||||
features = FeatureService.get_features(current_tenant_id, exclude_vector_space=True)
|
||||
if features.billing.enabled:
|
||||
members = features.members
|
||||
apps = features.apps
|
||||
vector_space = features.vector_space
|
||||
documents_upload_quota = features.documents_upload_quota
|
||||
annotation_quota_limit = features.annotation_quota_limit
|
||||
if resource == "members" and 0 < members.limit <= members.size:
|
||||
abort(403, "The number of members has reached the limit of your subscription.")
|
||||
elif resource == "apps" and 0 < apps.limit <= apps.size:
|
||||
abort(403, "The number of apps has reached the limit of your subscription.")
|
||||
elif resource == "vector_space" and 0 < vector_space.limit <= vector_space.size:
|
||||
abort(
|
||||
403, "The capacity of the knowledge storage space has reached the limit of your subscription."
|
||||
)
|
||||
elif resource == "documents" and 0 < documents_upload_quota.limit <= documents_upload_quota.size:
|
||||
# The api of file upload is used in the multiple places,
|
||||
# so we need to check the source of the request from datasets
|
||||
@ -140,7 +147,7 @@ def cloud_edition_billing_knowledge_limit_check[**P, R](
|
||||
@wraps(view)
|
||||
def decorated(*args: P.args, **kwargs: P.kwargs):
|
||||
_, current_tenant_id = current_account_with_tenant()
|
||||
features = FeatureService.get_features(current_tenant_id)
|
||||
features = FeatureService.get_features(current_tenant_id, exclude_vector_space=True)
|
||||
if features.billing.enabled:
|
||||
if resource == "add_segment":
|
||||
if features.billing.subscription.plan == CloudPlan.SANDBOX:
|
||||
@ -291,7 +298,7 @@ def knowledge_pipeline_publish_enabled[**P, R](view: Callable[P, R]) -> Callable
|
||||
@wraps(view)
|
||||
def decorated(*args: P.args, **kwargs: P.kwargs):
|
||||
_, current_tenant_id = current_account_with_tenant()
|
||||
features = FeatureService.get_features(current_tenant_id)
|
||||
features = FeatureService.get_features(current_tenant_id, exclude_vector_space=True)
|
||||
if features.knowledge_pipeline.publish_enabled:
|
||||
return view(*args, **kwargs)
|
||||
abort(403)
|
||||
|
||||
@ -13,6 +13,7 @@ from pydantic import BaseModel
|
||||
from sqlalchemy import select
|
||||
from werkzeug.exceptions import Forbidden, NotFound, Unauthorized
|
||||
|
||||
from configs import dify_config
|
||||
from enums.cloud_plan import CloudPlan
|
||||
from extensions.ext_database import db
|
||||
from extensions.ext_redis import redis_client
|
||||
@ -140,20 +141,26 @@ def cloud_edition_billing_resource_check[**P, R](
|
||||
def interceptor(view: Callable[P, R]):
|
||||
def decorated(*args: P.args, **kwargs: P.kwargs):
|
||||
api_token = validate_and_get_api_token(api_token_type)
|
||||
features = FeatureService.get_features(api_token.tenant_id)
|
||||
if resource == "vector_space":
|
||||
if not dify_config.BILLING_ENABLED:
|
||||
return view(*args, **kwargs)
|
||||
|
||||
vector_space = FeatureService.get_vector_space(api_token.tenant_id)
|
||||
if 0 < vector_space.limit <= vector_space.size:
|
||||
raise Forbidden("The capacity of the vector space has reached the limit of your subscription.")
|
||||
return view(*args, **kwargs)
|
||||
|
||||
features = FeatureService.get_features(api_token.tenant_id, exclude_vector_space=True)
|
||||
|
||||
if features.billing.enabled:
|
||||
members = features.members
|
||||
apps = features.apps
|
||||
vector_space = features.vector_space
|
||||
documents_upload_quota = features.documents_upload_quota
|
||||
|
||||
if resource == "members" and 0 < members.limit <= members.size:
|
||||
raise Forbidden("The number of members has reached the limit of your subscription.")
|
||||
elif resource == "apps" and 0 < apps.limit <= apps.size:
|
||||
raise Forbidden("The number of apps has reached the limit of your subscription.")
|
||||
elif resource == "vector_space" and 0 < vector_space.limit <= vector_space.size:
|
||||
raise Forbidden("The capacity of the vector space has reached the limit of your subscription.")
|
||||
elif resource == "documents" and 0 < documents_upload_quota.limit <= documents_upload_quota.size:
|
||||
raise Forbidden("The number of documents has reached the limit of your subscription.")
|
||||
else:
|
||||
@ -174,7 +181,7 @@ def cloud_edition_billing_knowledge_limit_check[**P, R](
|
||||
@wraps(view)
|
||||
def decorated(*args: P.args, **kwargs: P.kwargs):
|
||||
api_token = validate_and_get_api_token(api_token_type)
|
||||
features = FeatureService.get_features(api_token.tenant_id)
|
||||
features = FeatureService.get_features(api_token.tenant_id, exclude_vector_space=True)
|
||||
if features.billing.enabled:
|
||||
if resource == "add_segment":
|
||||
if features.billing.subscription.plan == CloudPlan.SANDBOX:
|
||||
|
||||
@ -81,7 +81,7 @@ class AppSiteApi(WebApiResource):
|
||||
if app_model.tenant.status == TenantStatus.ARCHIVE:
|
||||
raise Forbidden()
|
||||
|
||||
can_replace_logo = FeatureService.get_features(app_model.tenant_id).can_replace_logo
|
||||
can_replace_logo = FeatureService.get_features(app_model.tenant_id, exclude_vector_space=True).can_replace_logo
|
||||
|
||||
return AppSiteInfo(app_model.tenant, app_model, site, end_user.id, can_replace_logo)
|
||||
|
||||
@ -119,6 +119,6 @@ def serialize_site(site: Site) -> dict[str, Any]:
|
||||
|
||||
|
||||
def serialize_app_site_payload(app_model: App, site: Site, end_user_id: str | None) -> dict[str, Any]:
|
||||
can_replace_logo = FeatureService.get_features(app_model.tenant_id).can_replace_logo
|
||||
can_replace_logo = FeatureService.get_features(app_model.tenant_id, exclude_vector_space=True).can_replace_logo
|
||||
app_site_info = AppSiteInfo(app_model.tenant, app_model, site, end_user_id, can_replace_logo)
|
||||
return cast(dict[str, Any], marshal(app_site_info, AppSiteApi.app_fields))
|
||||
|
||||
@ -534,7 +534,9 @@ class ProviderManager:
|
||||
cache_key = f"tenant:{tenant_id}:model_load_balancing_enabled"
|
||||
cache_result = redis_client.get(cache_key)
|
||||
if cache_result is None:
|
||||
model_load_balancing_enabled = FeatureService.get_features(tenant_id).model_load_balancing_enabled
|
||||
model_load_balancing_enabled = FeatureService.get_features(
|
||||
tenant_id, exclude_vector_space=True
|
||||
).model_load_balancing_enabled
|
||||
redis_client.setex(cache_key, 120, str(model_load_balancing_enabled))
|
||||
else:
|
||||
cache_result = cache_result.decode("utf-8")
|
||||
|
||||
@ -58,7 +58,7 @@ def check_workspace_owner_transfer_permission(workspace_id: str) -> None:
|
||||
Raises:
|
||||
Forbidden: If either billing plan or workspace policy prohibits ownership transfer
|
||||
"""
|
||||
features = FeatureService.get_features(workspace_id)
|
||||
features = FeatureService.get_features(workspace_id, exclude_vector_space=True)
|
||||
if not features.is_allow_transfer_workspace:
|
||||
raise Forbidden("Your current plan does not allow workspace ownership transfer")
|
||||
|
||||
|
||||
@ -112,7 +112,7 @@ def clean_unused_datasets_task():
|
||||
features_cache_key = f"features:{dataset.tenant_id}"
|
||||
plan_cache = redis_client.get(features_cache_key)
|
||||
if plan_cache is None:
|
||||
features = FeatureService.get_features(dataset.tenant_id)
|
||||
features = FeatureService.get_features(dataset.tenant_id, exclude_vector_space=True)
|
||||
redis_client.setex(features_cache_key, 600, features.billing.subscription.plan)
|
||||
plan = features.billing.subscription.plan
|
||||
else:
|
||||
|
||||
@ -45,7 +45,7 @@ def mail_clean_document_notify_task():
|
||||
dataset_auto_disable_logs_map[dataset_auto_disable_log.tenant_id].append(dataset_auto_disable_log)
|
||||
url = f"{dify_config.CONSOLE_WEB_URL}/datasets"
|
||||
for tenant_id, tenant_dataset_auto_disable_logs in dataset_auto_disable_logs_map.items():
|
||||
features = FeatureService.get_features(tenant_id)
|
||||
features = FeatureService.get_features(tenant_id, exclude_vector_space=True)
|
||||
plan = features.billing.subscription.plan
|
||||
if plan != CloudPlan.SANDBOX:
|
||||
knowledge_details = []
|
||||
|
||||
@ -521,7 +521,7 @@ class AppAnnotationService:
|
||||
)
|
||||
|
||||
# Check annotation quota limit
|
||||
features = FeatureService.get_features(current_tenant_id)
|
||||
features = FeatureService.get_features(current_tenant_id, exclude_vector_space=True)
|
||||
if features.billing.enabled:
|
||||
annotation_quota_limit = features.annotation_quota_limit
|
||||
if annotation_quota_limit.limit < len(result) + annotation_quota_limit.size:
|
||||
|
||||
@ -192,6 +192,9 @@ class BillingService:
|
||||
params["exclude_vector_space"] = "true"
|
||||
|
||||
billing_info = cls._send_request("GET", "/subscription/info", params=params)
|
||||
if exclude_vector_space and billing_info.get("vector_space") is None:
|
||||
# Unset proto message fields can be serialized as null; the light billing contract treats it as absent.
|
||||
billing_info.pop("vector_space", None)
|
||||
return _billing_info_adapter.validate_python(billing_info)
|
||||
|
||||
@classmethod
|
||||
|
||||
@ -1325,7 +1325,7 @@ class DatasetService:
|
||||
def get_dataset_auto_disable_logs(dataset_id: str) -> AutoDisableLogsDict:
|
||||
assert isinstance(current_user, Account)
|
||||
assert current_user.current_tenant_id is not None
|
||||
features = FeatureService.get_features(current_user.current_tenant_id)
|
||||
features = FeatureService.get_features(current_user.current_tenant_id, exclude_vector_space=True)
|
||||
if not features.billing.enabled or features.billing.subscription.plan == CloudPlan.SANDBOX:
|
||||
return {
|
||||
"document_ids": [],
|
||||
@ -2007,7 +2007,7 @@ class DocumentService:
|
||||
assert isinstance(current_user, Account)
|
||||
assert current_user.current_tenant_id is not None
|
||||
|
||||
features = FeatureService.get_features(current_user.current_tenant_id)
|
||||
features = FeatureService.get_features(current_user.current_tenant_id, exclude_vector_space=True)
|
||||
|
||||
if features.billing.enabled:
|
||||
if not knowledge_config.original_document_id:
|
||||
@ -2798,7 +2798,7 @@ class DocumentService:
|
||||
assert current_user.current_tenant_id is not None
|
||||
assert knowledge_config.data_source
|
||||
|
||||
features = FeatureService.get_features(current_user.current_tenant_id)
|
||||
features = FeatureService.get_features(current_user.current_tenant_id, exclude_vector_space=True)
|
||||
|
||||
if features.billing.enabled:
|
||||
count = 0
|
||||
|
||||
@ -41,7 +41,7 @@ class DocumentTaskProxyBase(ABC):
|
||||
|
||||
@cached_property
|
||||
def features(self):
|
||||
return FeatureService.get_features(self._tenant_id)
|
||||
return FeatureService.get_features(self._tenant_id, exclude_vector_space=True)
|
||||
|
||||
@abstractmethod
|
||||
def _send_to_direct_queue(self, task_func: Callable[..., Any]):
|
||||
|
||||
@ -130,7 +130,7 @@ class FeatureModel(FeatureResponseModel):
|
||||
education: EducationModel = EducationModel()
|
||||
members: LimitationModel = LimitationModel(size=0, limit=1)
|
||||
apps: LimitationModel = LimitationModel(size=0, limit=10)
|
||||
vector_space: LimitationModel = LimitationModel(size=0, limit=5)
|
||||
vector_space: LimitationModel | None = LimitationModel(size=0, limit=5)
|
||||
knowledge_rate_limit: int = 10
|
||||
annotation_quota_limit: LimitationModel = LimitationModel(size=0, limit=10)
|
||||
documents_upload_quota: LimitationModel = LimitationModel(size=0, limit=50)
|
||||
@ -188,6 +188,8 @@ class FeatureService:
|
||||
@classmethod
|
||||
def get_features(cls, tenant_id: str, exclude_vector_space: bool = False) -> FeatureModel:
|
||||
features = FeatureModel()
|
||||
if exclude_vector_space:
|
||||
features.vector_space = None
|
||||
|
||||
cls._fulfill_params_from_env(features)
|
||||
|
||||
@ -347,6 +349,7 @@ class FeatureService:
|
||||
features.apps.limit = billing_info["apps"]["limit"]
|
||||
|
||||
if not exclude_vector_space:
|
||||
assert features.vector_space is not None
|
||||
cls._fulfill_vector_space_from_billing_info(features.vector_space, billing_info)
|
||||
|
||||
if "documents_upload_quota" in billing_info:
|
||||
|
||||
@ -136,7 +136,7 @@ class EmailDeliveryTestHandler:
|
||||
) -> DeliveryTestResult:
|
||||
if not isinstance(method, EmailDeliveryMethod):
|
||||
raise DeliveryTestUnsupportedError("Delivery method does not support test send.")
|
||||
features = FeatureService.get_features(context.tenant_id)
|
||||
features = FeatureService.get_features(context.tenant_id, exclude_vector_space=True)
|
||||
if not features.human_input_email_delivery_enabled:
|
||||
raise DeliveryTestError("Email delivery is not available for current plan.")
|
||||
if not mail.is_inited():
|
||||
|
||||
@ -29,7 +29,7 @@ class RagPipelineTaskProxy:
|
||||
|
||||
@cached_property
|
||||
def features(self):
|
||||
return FeatureService.get_features(self._dataset_tenant_id)
|
||||
return FeatureService.get_features(self._dataset_tenant_id, exclude_vector_space=True)
|
||||
|
||||
def _upload_invoke_entities(self) -> str:
|
||||
text = [item.model_dump() for item in self._rag_pipeline_invoke_entities]
|
||||
|
||||
@ -33,7 +33,7 @@ class WorkspaceService:
|
||||
assert tenant_account_join is not None, "TenantAccountJoin not found"
|
||||
tenant_info["role"] = tenant_account_join.role
|
||||
|
||||
feature = FeatureService.get_features(tenant.id)
|
||||
feature = FeatureService.get_features(tenant.id, exclude_vector_space=True)
|
||||
can_replace_logo = feature.can_replace_logo
|
||||
|
||||
if can_replace_logo and TenantService.has_roles(tenant, [TenantAccountRole.OWNER, TenantAccountRole.ADMIN]):
|
||||
|
||||
@ -66,6 +66,7 @@ def _document_indexing(dataset_id: str, document_ids: Sequence[str]):
|
||||
try:
|
||||
if features.billing.enabled:
|
||||
vector_space = features.vector_space
|
||||
assert vector_space is not None
|
||||
count = len(document_ids)
|
||||
batch_upload_limit = int(dify_config.BATCH_UPLOAD_LIMIT)
|
||||
if features.billing.subscription.plan == CloudPlan.SANDBOX and count > 1:
|
||||
|
||||
@ -92,6 +92,7 @@ def _duplicate_document_indexing_task(dataset_id: str, document_ids: Sequence[st
|
||||
try:
|
||||
if features.billing.enabled:
|
||||
vector_space = features.vector_space
|
||||
assert vector_space is not None
|
||||
count = len(document_ids)
|
||||
if features.billing.subscription.plan == CloudPlan.SANDBOX and count > 1:
|
||||
raise ValueError("Your current plan does not support batch upload, please upgrade your plan.")
|
||||
|
||||
@ -157,7 +157,7 @@ def dispatch_human_input_email_task(form_id: str, node_title: str | None = None,
|
||||
if form is None:
|
||||
logger.warning("Human input form not found, form_id=%s", form_id)
|
||||
return
|
||||
features = FeatureService.get_features(form.tenant_id)
|
||||
features = FeatureService.get_features(form.tenant_id, exclude_vector_space=True)
|
||||
if not features.human_input_email_delivery_enabled:
|
||||
logger.info(
|
||||
"Human input email delivery is not available for tenant=%s, form_id=%s",
|
||||
|
||||
@ -52,6 +52,7 @@ def retry_document_indexing_task(dataset_id: str, document_ids: list[str], user_
|
||||
try:
|
||||
if features.billing.enabled:
|
||||
vector_space = features.vector_space
|
||||
assert vector_space is not None
|
||||
if 0 < vector_space.limit <= vector_space.size:
|
||||
raise ValueError(
|
||||
"Your total number of documents plus the number of uploads have over the limit of "
|
||||
|
||||
@ -39,6 +39,7 @@ def sync_website_document_indexing_task(dataset_id: str, document_id: str):
|
||||
try:
|
||||
if features.billing.enabled:
|
||||
vector_space = features.vector_space
|
||||
assert vector_space is not None
|
||||
if 0 < vector_space.limit <= vector_space.size:
|
||||
raise ValueError(
|
||||
"Your total number of documents plus the number of uploads have over the limit of "
|
||||
|
||||
@ -127,7 +127,7 @@ class TestEmailDeliveryTestHandler:
|
||||
monkeypatch.setattr(
|
||||
service_module.FeatureService,
|
||||
"get_features",
|
||||
lambda _tenant_id: SimpleNamespace(human_input_email_delivery_enabled=False),
|
||||
lambda _tenant_id, **_kwargs: SimpleNamespace(human_input_email_delivery_enabled=False),
|
||||
)
|
||||
handler = EmailDeliveryTestHandler(session_factory=MagicMock())
|
||||
context = DeliveryTestContext(
|
||||
@ -142,7 +142,7 @@ class TestEmailDeliveryTestHandler:
|
||||
monkeypatch.setattr(
|
||||
service_module.FeatureService,
|
||||
"get_features",
|
||||
lambda _id: SimpleNamespace(human_input_email_delivery_enabled=True),
|
||||
lambda _id, **_kwargs: SimpleNamespace(human_input_email_delivery_enabled=True),
|
||||
)
|
||||
monkeypatch.setattr(service_module.mail, "is_inited", lambda: False)
|
||||
|
||||
@ -159,7 +159,7 @@ class TestEmailDeliveryTestHandler:
|
||||
monkeypatch.setattr(
|
||||
service_module.FeatureService,
|
||||
"get_features",
|
||||
lambda _id: SimpleNamespace(human_input_email_delivery_enabled=True),
|
||||
lambda _id, **_kwargs: SimpleNamespace(human_input_email_delivery_enabled=True),
|
||||
)
|
||||
monkeypatch.setattr(service_module.mail, "is_inited", lambda: True)
|
||||
|
||||
@ -178,7 +178,7 @@ class TestEmailDeliveryTestHandler:
|
||||
monkeypatch.setattr(
|
||||
service_module.FeatureService,
|
||||
"get_features",
|
||||
lambda _id: SimpleNamespace(human_input_email_delivery_enabled=True),
|
||||
lambda _id, **_kwargs: SimpleNamespace(human_input_email_delivery_enabled=True),
|
||||
)
|
||||
monkeypatch.setattr(service_module.mail, "is_inited", lambda: True)
|
||||
mock_mail_send = MagicMock()
|
||||
@ -214,7 +214,7 @@ class TestEmailDeliveryTestHandler:
|
||||
monkeypatch.setattr(
|
||||
service_module.FeatureService,
|
||||
"get_features",
|
||||
lambda _id: SimpleNamespace(human_input_email_delivery_enabled=True),
|
||||
lambda _id, **_kwargs: SimpleNamespace(human_input_email_delivery_enabled=True),
|
||||
)
|
||||
monkeypatch.setattr(service_module.mail, "is_inited", lambda: True)
|
||||
mock_mail_send = MagicMock()
|
||||
|
||||
@ -255,11 +255,43 @@ class TestBillingResourceLimits:
|
||||
with patch(
|
||||
"controllers.console.wraps.current_account_with_tenant", return_value=(MockUser("test_user"), "tenant123")
|
||||
):
|
||||
with patch("controllers.console.wraps.FeatureService.get_features", return_value=mock_features):
|
||||
with patch(
|
||||
"controllers.console.wraps.FeatureService.get_features", return_value=mock_features
|
||||
) as get_features:
|
||||
result = add_member()
|
||||
|
||||
# Assert
|
||||
assert result == "member_added"
|
||||
get_features.assert_called_once_with("tenant123", exclude_vector_space=True)
|
||||
|
||||
def test_should_load_vector_space_from_dedicated_quota_api(self):
|
||||
"""Test vector-space limit checks avoid loading the full feature payload."""
|
||||
# Arrange
|
||||
mock_vector_space = MagicMock()
|
||||
mock_vector_space.limit = 10
|
||||
mock_vector_space.size = 5
|
||||
|
||||
@cloud_edition_billing_resource_check("vector_space")
|
||||
def add_segment():
|
||||
return "segment_added"
|
||||
|
||||
# Act
|
||||
with patch(
|
||||
"controllers.console.wraps.current_account_with_tenant", return_value=(MockUser("test_user"), "tenant123")
|
||||
):
|
||||
with (
|
||||
patch("controllers.console.wraps.dify_config.BILLING_ENABLED", True),
|
||||
patch(
|
||||
"controllers.console.wraps.FeatureService.get_vector_space", return_value=mock_vector_space
|
||||
) as get_vector_space,
|
||||
patch("controllers.console.wraps.FeatureService.get_features") as get_features,
|
||||
):
|
||||
result = add_segment()
|
||||
|
||||
# Assert
|
||||
assert result == "segment_added"
|
||||
get_vector_space.assert_called_once_with("tenant123")
|
||||
get_features.assert_not_called()
|
||||
|
||||
def test_should_reject_when_over_resource_limit(self):
|
||||
"""Test that requests are rejected when over resource limits"""
|
||||
|
||||
@ -139,7 +139,7 @@ class TestTenantListApi:
|
||||
assert result["workspaces"][0]["plan"] == CloudPlan.TEAM
|
||||
assert result["workspaces"][1]["plan"] == CloudPlan.PROFESSIONAL
|
||||
get_plan_bulk_mock.assert_called_once_with(["t1", "t2"])
|
||||
get_features_mock.assert_called_once_with("t2")
|
||||
get_features_mock.assert_called_once_with("t2", exclude_vector_space=True)
|
||||
|
||||
def test_get_saas_path_falls_back_to_legacy_feature_path_on_bulk_error(self, app: Flask):
|
||||
"""Test fallback to FeatureService when bulk billing returns empty result.
|
||||
@ -235,7 +235,7 @@ class TestTenantListApi:
|
||||
|
||||
assert status == 200
|
||||
assert result["workspaces"][0]["plan"] == CloudPlan.SANDBOX
|
||||
get_features_mock.assert_called_once_with("t1")
|
||||
get_features_mock.assert_called_once_with("t1", exclude_vector_space=True)
|
||||
|
||||
def test_get_enterprise_only_skips_feature_service(self, app: Flask):
|
||||
api = TenantListApi()
|
||||
|
||||
@ -872,6 +872,11 @@ class TestSegmentApiPost:
|
||||
mock_features.billing.enabled = False
|
||||
mock_feature_svc.get_features.return_value = mock_features
|
||||
|
||||
mock_vector_space = Mock()
|
||||
mock_vector_space.limit = 10
|
||||
mock_vector_space.size = 0
|
||||
mock_feature_svc.get_vector_space.return_value = mock_vector_space
|
||||
|
||||
mock_rate_limit = Mock()
|
||||
mock_rate_limit.enabled = False
|
||||
mock_feature_svc.get_knowledge_rate_limit.return_value = mock_rate_limit
|
||||
@ -1209,6 +1214,10 @@ class TestDatasetSegmentApiUpdate:
|
||||
mock_features = Mock()
|
||||
mock_features.billing.enabled = False
|
||||
mock_feature_svc.get_features.return_value = mock_features
|
||||
mock_vector_space = Mock()
|
||||
mock_vector_space.limit = 10
|
||||
mock_vector_space.size = 0
|
||||
mock_feature_svc.get_vector_space.return_value = mock_vector_space
|
||||
mock_rate_limit = Mock()
|
||||
mock_rate_limit.enabled = False
|
||||
mock_feature_svc.get_knowledge_rate_limit.return_value = mock_rate_limit
|
||||
@ -1710,6 +1719,10 @@ class TestChildChunkApiPost:
|
||||
mock_features = Mock()
|
||||
mock_features.billing.enabled = False
|
||||
mock_feature_svc.get_features.return_value = mock_features
|
||||
mock_vector_space = Mock()
|
||||
mock_vector_space.limit = 10
|
||||
mock_vector_space.size = 0
|
||||
mock_feature_svc.get_vector_space.return_value = mock_vector_space
|
||||
mock_rate_limit = Mock()
|
||||
mock_rate_limit.enabled = False
|
||||
mock_feature_svc.get_knowledge_rate_limit.return_value = mock_rate_limit
|
||||
|
||||
@ -950,7 +950,8 @@ class TestDocumentAddByTextApi:
|
||||
"""Configure mocks to neutralise billing/auth decorators.
|
||||
|
||||
``cloud_edition_billing_resource_check`` calls
|
||||
``FeatureService.get_features`` and
|
||||
``FeatureService.get_vector_space`` for vector-space checks and
|
||||
``FeatureService.get_features`` for other resource checks.
|
||||
``cloud_edition_billing_rate_limit_check`` calls
|
||||
``FeatureService.get_knowledge_rate_limit``.
|
||||
Both call ``validate_and_get_api_token`` first.
|
||||
@ -963,6 +964,11 @@ class TestDocumentAddByTextApi:
|
||||
mock_features.billing.enabled = False
|
||||
mock_feature_svc.get_features.return_value = mock_features
|
||||
|
||||
mock_vector_space = Mock()
|
||||
mock_vector_space.limit = 10
|
||||
mock_vector_space.size = 0
|
||||
mock_feature_svc.get_vector_space.return_value = mock_vector_space
|
||||
|
||||
mock_rate_limit = Mock()
|
||||
mock_rate_limit.enabled = False
|
||||
mock_feature_svc.get_knowledge_rate_limit.return_value = mock_rate_limit
|
||||
@ -1140,6 +1146,10 @@ def _setup_billing_mocks(mock_validate_token, mock_feature_svc, tenant_id: str):
|
||||
mock_features = Mock()
|
||||
mock_features.billing.enabled = False
|
||||
mock_feature_svc.get_features.return_value = mock_features
|
||||
mock_vector_space = Mock()
|
||||
mock_vector_space.limit = 10
|
||||
mock_vector_space.size = 0
|
||||
mock_feature_svc.get_vector_space.return_value = mock_vector_space
|
||||
mock_rate_limit = Mock()
|
||||
mock_rate_limit.enabled = False
|
||||
mock_feature_svc.get_knowledge_rate_limit.return_value = mock_rate_limit
|
||||
|
||||
@ -265,6 +265,65 @@ class TestCloudEditionBillingResourceCheck:
|
||||
|
||||
# Assert
|
||||
assert result == "member_added"
|
||||
mock_get_features.assert_called_once_with("tenant123", exclude_vector_space=True)
|
||||
|
||||
@patch("controllers.service_api.wraps.validate_and_get_api_token")
|
||||
@patch("controllers.service_api.wraps.FeatureService.get_features")
|
||||
@patch("controllers.service_api.wraps.FeatureService.get_vector_space")
|
||||
def test_loads_vector_space_from_dedicated_quota_api(
|
||||
self, mock_get_vector_space, mock_get_features, mock_validate_token, app: Flask
|
||||
):
|
||||
"""Test vector-space resource checks avoid loading the full feature payload."""
|
||||
# Arrange
|
||||
mock_validate_token.return_value = Mock(tenant_id="tenant123")
|
||||
|
||||
mock_vector_space = Mock()
|
||||
mock_vector_space.limit = 10
|
||||
mock_vector_space.size = 5
|
||||
mock_get_vector_space.return_value = mock_vector_space
|
||||
|
||||
@cloud_edition_billing_resource_check("vector_space", "dataset")
|
||||
def add_segment():
|
||||
return "segment_added"
|
||||
|
||||
# Act
|
||||
with (
|
||||
app.test_request_context("/", method="GET"),
|
||||
patch("controllers.service_api.wraps.dify_config.BILLING_ENABLED", True),
|
||||
):
|
||||
result = add_segment()
|
||||
|
||||
# Assert
|
||||
assert result == "segment_added"
|
||||
mock_get_vector_space.assert_called_once_with("tenant123")
|
||||
mock_get_features.assert_not_called()
|
||||
|
||||
@patch("controllers.service_api.wraps.validate_and_get_api_token")
|
||||
@patch("controllers.service_api.wraps.FeatureService.get_features")
|
||||
def test_loads_features_when_checking_non_vector_space_limit(
|
||||
self, mock_get_features, mock_validate_token, app: Flask
|
||||
):
|
||||
"""Test non-vector-space resource checks keep using the light feature payload."""
|
||||
# Arrange
|
||||
mock_validate_token.return_value = Mock(tenant_id="tenant123")
|
||||
|
||||
mock_features = Mock()
|
||||
mock_features.billing.enabled = True
|
||||
mock_features.documents_upload_quota.limit = 10
|
||||
mock_features.documents_upload_quota.size = 5
|
||||
mock_get_features.return_value = mock_features
|
||||
|
||||
@cloud_edition_billing_resource_check("documents", "dataset")
|
||||
def upload_document():
|
||||
return "document_uploaded"
|
||||
|
||||
# Act
|
||||
with app.test_request_context("/", method="GET"):
|
||||
result = upload_document()
|
||||
|
||||
# Assert
|
||||
assert result == "document_uploaded"
|
||||
mock_get_features.assert_called_once_with("tenant123", exclude_vector_space=True)
|
||||
|
||||
@patch("controllers.service_api.wraps.validate_and_get_api_token")
|
||||
@patch("controllers.service_api.wraps.FeatureService.get_features")
|
||||
|
||||
@ -126,7 +126,7 @@ def test_get_form_includes_site(monkeypatch: pytest.MonkeyPatch, app: Flask):
|
||||
monkeypatch.setattr(
|
||||
site_module.FeatureService,
|
||||
"get_features",
|
||||
lambda tenant_id: SimpleNamespace(can_replace_logo=True),
|
||||
lambda tenant_id, **_kwargs: SimpleNamespace(can_replace_logo=True),
|
||||
)
|
||||
|
||||
with app.test_request_context("/api/form/human_input/token-1", method="GET"):
|
||||
@ -245,7 +245,7 @@ def test_get_form_allows_backstage_token(monkeypatch: pytest.MonkeyPatch, app: F
|
||||
monkeypatch.setattr(
|
||||
site_module.FeatureService,
|
||||
"get_features",
|
||||
lambda tenant_id: SimpleNamespace(can_replace_logo=True),
|
||||
lambda tenant_id, **_kwargs: SimpleNamespace(can_replace_logo=True),
|
||||
)
|
||||
|
||||
with app.test_request_context("/api/form/human_input/token-1", method="GET"):
|
||||
|
||||
@ -36,7 +36,7 @@ class TestWorkspacePermissionHelper:
|
||||
# Should not raise
|
||||
check_workspace_owner_transfer_permission("test-workspace-id")
|
||||
|
||||
mock_feature_service.get_features.assert_called_once_with("test-workspace-id")
|
||||
mock_feature_service.get_features.assert_called_once_with("test-workspace-id", exclude_vector_space=True)
|
||||
|
||||
@patch("libs.workspace_permission.EnterpriseService")
|
||||
@patch("libs.workspace_permission.dify_config")
|
||||
|
||||
@ -344,7 +344,7 @@ class TestDispatchRouting:
|
||||
proxy._dispatch()
|
||||
|
||||
# Assert
|
||||
mock_features.assert_called_once_with(TENANT_ID)
|
||||
mock_features.assert_called_once_with(TENANT_ID, exclude_vector_space=True)
|
||||
|
||||
|
||||
class TestBaseRouterHelpers:
|
||||
|
||||
@ -343,6 +343,37 @@ class TestBillingServiceSubscriptionInfo:
|
||||
params={"tenant_id": tenant_id, "exclude_vector_space": "true"},
|
||||
)
|
||||
|
||||
def test_get_info_exclude_vector_space_normalizes_null_field(self, mock_send_request):
|
||||
"""When billing serializes skipped vector_space as null, get_info treats it as absent."""
|
||||
# Arrange
|
||||
tenant_id = "tenant-123"
|
||||
expected_response = {
|
||||
"enabled": True,
|
||||
"subscription": {"plan": "professional", "interval": "month", "education": False},
|
||||
"members": {"size": 1, "limit": 50},
|
||||
"apps": {"size": 1, "limit": 200},
|
||||
"vector_space": None,
|
||||
"knowledge_rate_limit": {"limit": 1000},
|
||||
"documents_upload_quota": {"size": 0, "limit": 1000},
|
||||
"annotation_quota_limit": {"size": 0, "limit": 5000},
|
||||
"docs_processing": "top-priority",
|
||||
"can_replace_logo": True,
|
||||
"model_load_balancing_enabled": True,
|
||||
"knowledge_pipeline_publish_enabled": True,
|
||||
}
|
||||
mock_send_request.return_value = expected_response
|
||||
|
||||
# Act
|
||||
result = BillingService.get_info(tenant_id, exclude_vector_space=True)
|
||||
|
||||
# Assert
|
||||
assert "vector_space" not in result
|
||||
mock_send_request.assert_called_once_with(
|
||||
"GET",
|
||||
"/subscription/info",
|
||||
params={"tenant_id": tenant_id, "exclude_vector_space": "true"},
|
||||
)
|
||||
|
||||
def test_get_vector_space_success(self, mock_send_request):
|
||||
"""Test successful retrieval of vector-space usage and limit."""
|
||||
# Arrange
|
||||
|
||||
@ -39,7 +39,7 @@ def fake_features(monkeypatch: pytest.MonkeyPatch):
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"services.dataset_service.FeatureService.get_features",
|
||||
lambda tenant_id: features,
|
||||
lambda tenant_id, **_kwargs: features,
|
||||
)
|
||||
return features
|
||||
|
||||
|
||||
@ -75,7 +75,7 @@ class TestDocumentIndexingTaskProxy:
|
||||
assert features1 == mock_features
|
||||
assert features2 == mock_features
|
||||
assert features1 is features2 # Should be the same instance due to caching
|
||||
mock_feature_service.get_features.assert_called_once_with("tenant-123")
|
||||
mock_feature_service.get_features.assert_called_once_with("tenant-123", exclude_vector_space=True)
|
||||
|
||||
@patch("services.document_indexing_proxy.document_indexing_task_proxy.normal_document_indexing_task")
|
||||
def test_send_to_direct_queue(self, mock_task):
|
||||
|
||||
@ -94,7 +94,7 @@ class TestDuplicateDocumentIndexingTaskProxy:
|
||||
assert features1 == mock_features
|
||||
assert features2 == mock_features
|
||||
assert features1 is features2 # Should be the same instance due to caching
|
||||
mock_feature_service.get_features.assert_called_once_with("tenant-123")
|
||||
mock_feature_service.get_features.assert_called_once_with("tenant-123", exclude_vector_space=True)
|
||||
|
||||
@patch(
|
||||
"services.document_indexing_proxy.duplicate_document_indexing_task_proxy.normal_duplicate_document_indexing_task"
|
||||
|
||||
@ -0,0 +1,37 @@
|
||||
from unittest.mock import patch
|
||||
|
||||
from services.feature_service import FeatureService
|
||||
|
||||
|
||||
def test_get_features_exclude_vector_space_sets_vector_space_to_none():
|
||||
tenant_id = "tenant-id"
|
||||
billing_info = {
|
||||
"enabled": True,
|
||||
"subscription": {"plan": "pro", "interval": "monthly", "education": False},
|
||||
"members": {"size": 1, "limit": 10},
|
||||
"apps": {"size": 2, "limit": 20},
|
||||
"documents_upload_quota": {"size": 3, "limit": 100},
|
||||
"annotation_quota_limit": {"size": 4, "limit": 50},
|
||||
"docs_processing": "standard",
|
||||
"can_replace_logo": True,
|
||||
"model_load_balancing_enabled": True,
|
||||
"knowledge_rate_limit": {"limit": 100},
|
||||
"knowledge_pipeline_publish_enabled": True,
|
||||
}
|
||||
|
||||
with (
|
||||
patch("services.feature_service.dify_config") as mock_config,
|
||||
patch("services.feature_service.BillingService.get_info", return_value=billing_info) as get_info,
|
||||
patch("services.feature_service.BillingService.get_quota_info", return_value={}),
|
||||
):
|
||||
mock_config.BILLING_ENABLED = True
|
||||
mock_config.ENTERPRISE_ENABLED = False
|
||||
mock_config.CAN_REPLACE_LOGO = False
|
||||
mock_config.MODEL_LB_ENABLED = False
|
||||
mock_config.DATASET_OPERATOR_ENABLED = False
|
||||
mock_config.EDUCATION_ENABLED = False
|
||||
|
||||
features = FeatureService.get_features(tenant_id, exclude_vector_space=True)
|
||||
|
||||
assert features.vector_space is None
|
||||
get_info.assert_called_once_with(tenant_id, exclude_vector_space=True)
|
||||
@ -144,7 +144,7 @@ class TestRagPipelineTaskProxy:
|
||||
assert features1 == mock_features
|
||||
assert features2 == mock_features
|
||||
assert features1 is features2 # Should be the same instance due to caching
|
||||
mock_feature_service.get_features.assert_called_once_with("tenant-123")
|
||||
mock_feature_service.get_features.assert_called_once_with("tenant-123", exclude_vector_space=True)
|
||||
|
||||
@patch("services.rag_pipeline.rag_pipeline_task_proxy.FileService")
|
||||
@patch("services.rag_pipeline.rag_pipeline_task_proxy.db")
|
||||
|
||||
@ -54,7 +54,7 @@ def test_dispatch_human_input_email_task_sends_to_each_recipient(monkeypatch: py
|
||||
monkeypatch.setattr(
|
||||
task_module.FeatureService,
|
||||
"get_features",
|
||||
lambda _tenant_id: SimpleNamespace(human_input_email_delivery_enabled=True),
|
||||
lambda _tenant_id, **_kwargs: SimpleNamespace(human_input_email_delivery_enabled=True),
|
||||
)
|
||||
jobs: Sequence[task_module._EmailDeliveryJob] = [_build_job(recipient_count=2)]
|
||||
monkeypatch.setattr(task_module, "_load_email_jobs", lambda _session, _form: jobs)
|
||||
@ -78,7 +78,7 @@ def test_dispatch_human_input_email_task_skips_when_feature_disabled(monkeypatch
|
||||
monkeypatch.setattr(
|
||||
task_module.FeatureService,
|
||||
"get_features",
|
||||
lambda _tenant_id: SimpleNamespace(human_input_email_delivery_enabled=False),
|
||||
lambda _tenant_id, **_kwargs: SimpleNamespace(human_input_email_delivery_enabled=False),
|
||||
)
|
||||
monkeypatch.setattr(task_module, "_load_email_jobs", lambda _session, _form: [])
|
||||
|
||||
@ -109,7 +109,7 @@ def test_dispatch_human_input_email_task_replaces_body_variables(monkeypatch: py
|
||||
monkeypatch.setattr(
|
||||
task_module.FeatureService,
|
||||
"get_features",
|
||||
lambda _tenant_id: SimpleNamespace(human_input_email_delivery_enabled=True),
|
||||
lambda _tenant_id, **_kwargs: SimpleNamespace(human_input_email_delivery_enabled=True),
|
||||
)
|
||||
monkeypatch.setattr(task_module, "_load_email_jobs", lambda _session, _form: [job])
|
||||
monkeypatch.setattr(task_module, "_load_variable_pool", lambda _workflow_run_id: variable_pool)
|
||||
@ -142,7 +142,7 @@ def test_dispatch_human_input_email_task_sanitizes_subject(
|
||||
monkeypatch.setattr(
|
||||
task_module.FeatureService,
|
||||
"get_features",
|
||||
lambda _tenant_id: SimpleNamespace(human_input_email_delivery_enabled=True),
|
||||
lambda _tenant_id, **_kwargs: SimpleNamespace(human_input_email_delivery_enabled=True),
|
||||
)
|
||||
monkeypatch.setattr(task_module, "_load_email_jobs", lambda _session, _form: [job])
|
||||
monkeypatch.setattr(task_module, "_load_variable_pool", lambda _workflow_run_id: None)
|
||||
|
||||
Loading…
Reference in New Issue
Block a user