feat:精简

backend/open_webui/utils/subscription/check.py

@@ -10,7 +10,7 @@ from open_webui.models.subscriptions import (
 )
 
 
-async def check_subscription_access(
+def check_subscription_access(
     user_id: str,
     model_id: str,
     chat_id: Optional[str] = None,
@@ -21,8 +21,8 @@ async def check_subscription_access(
 
     Raises HTTPException if:
     1. User's subscription has expired (auto-downgrades to Free)
-    2. Monthly message limit reached
-    3. Model is not in allowed list
+    2. Model usage limit reached for this billing period
+    3. Model is not allowed (limit = 0)
     """
     now = int(time.time())
 
@@ -47,19 +47,23 @@ async def check_subscription_access(
             detail="Subscription plan not found"
         )
 
-    # check model access
-    if plan.allowed_models and model_id not in plan.allowed_models:
+    # get model limit for this plan
+    model_limit = plan.get_model_limit(model_id)
+
+    # check if model is not allowed (limit = 0)
+    if model_limit == 0:
         raise HTTPException(
             status_code=403,
-            detail=f"Your {plan.name} plan does not include access to this model. Please upgrade your subscription."
+            detail=f"您的 {plan.name} 套餐不支持使用此模型。请升级套餐以获取访问权限。"
         )
 
-    # check usage limit
-    if plan.monthly_message_limit is not None:
-        if subscription.messages_used >= plan.monthly_message_limit:
+    # check usage limit (skip if unlimited = -1)
+    if model_limit > 0:
+        current_usage = subscription.get_model_usage(model_id)
+        if current_usage >= model_limit:
             raise HTTPException(
                 status_code=403,
-                detail=f"You have reached your monthly message limit ({plan.monthly_message_limit}). Please wait until next month or upgrade your subscription."
+                detail=f"您本月已使用 {current_usage}/{model_limit} 次此模型。请等待下月刷新或升级套餐。"
             )
 
 
@@ -77,8 +81,11 @@ def record_usage(
     if not subscription:
         return
 
-    # increment usage counter
-    UserSubscriptions.increment_usage(subscription.id)
+    # increment per-model usage counter
+    if model_id:
+        UserSubscriptions.increment_model_usage(subscription.id, model_id)
+    else:
+        UserSubscriptions.increment_usage(subscription.id)
 
     # log usage
     SubscriptionUsageLogs.insert(
@@ -91,27 +98,100 @@ def record_usage(
     )
 
 
+def get_user_model_limit(user_id: str, model_id: str) -> int:
+    """
+    Get the usage limit for a specific model for a user.
+    Returns: -1 = unlimited, 0 = not allowed, positive = monthly limit
+    """
+    subscription = UserSubscriptions.get_by_user_id(user_id)
+    if not subscription:
+        default_plan = SubscriptionPlans.get_default_plan()
+        if default_plan:
+            return default_plan.get_model_limit(model_id)
+        return -1  # allow by default if no plan
+
+    plan = SubscriptionPlans.get_plan_by_id(subscription.plan_id)
+    if not plan:
+        return -1
+
+    return plan.get_model_limit(model_id)
+
+
+def get_user_model_remaining(user_id: str, model_id: str) -> Optional[int]:
+    """
+    Get the remaining usage count for a specific model for a user.
+    Returns: None if unlimited, 0 if not allowed or exhausted, positive if remaining
+    """
+    subscription = UserSubscriptions.get_by_user_id(user_id)
+    if not subscription:
+        default_plan = SubscriptionPlans.get_default_plan()
+        if default_plan:
+            limit = default_plan.get_model_limit(model_id)
+            return None if limit == -1 else limit
+        return None
+
+    plan = SubscriptionPlans.get_plan_by_id(subscription.plan_id)
+    if not plan:
+        return None
+
+    limit = plan.get_model_limit(model_id)
+    if limit == -1:
+        return None
+    if limit == 0:
+        return 0
+
+    used = subscription.get_model_usage(model_id)
+    return max(0, limit - used)
+
+
 def get_user_allowed_models(user_id: str) -> Optional[list[str]]:
     """
+    Deprecated: Use get_user_model_limit instead.
     Get the list of models allowed for a user based on their subscription.
     Returns None if all models are allowed.
     """
     subscription = UserSubscriptions.get_by_user_id(user_id)
     if not subscription:
         default_plan = SubscriptionPlans.get_default_plan()
-        return default_plan.allowed_models if default_plan else None
+        if default_plan and default_plan.model_limits:
+            # return models with limit != 0
+            return [m for m, l in default_plan.model_limits.items() if l != 0]
+        return None
 
     plan = SubscriptionPlans.get_plan_by_id(subscription.plan_id)
-    return plan.allowed_models if plan else None
+    if plan and plan.model_limits:
+        return [m for m, l in plan.model_limits.items() if l != 0]
+    return None
 
 
 def filter_models_by_subscription(user_id: str, models: list) -> list:
     """
     Filter a list of models based on user's subscription.
+    Models with limit = 0 are filtered out.
     """
-    allowed_models = get_user_allowed_models(user_id)
+    subscription = UserSubscriptions.get_by_user_id(user_id)
+    plan = None
 
-    if allowed_models is None:
+    if subscription:
+        plan = SubscriptionPlans.get_plan_by_id(subscription.plan_id)
+    else:
+        plan = SubscriptionPlans.get_default_plan()
+
+    if not plan:
         return models
 
-    return [m for m in models if m.get("id") in allowed_models]
+    # if no model_limits defined, use default_model_limit
+    if not plan.model_limits:
+        if plan.default_model_limit == 0:
+            return []  # no models allowed
+        return models  # all models allowed
+
+    # filter out models with limit = 0
+    result = []
+    for m in models:
+        model_id = m.get("id", "")
+        limit = plan.get_model_limit(model_id)
+        if limit != 0:
+            result.append(m)
+
+    return result